#!/usr/bin/env python
import socket
import sys
import os

# Unix socket
server_address = '/var/run/rrdcached.sock'

# TCP socket
#server_address = 'localhost:42217'

sock = None
try:
    if os.path.exists(server_address):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    else:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if ':' in server_address:
            split = server_address.rsplit(':', 1)
            server_address = (split[0],int(split[1]))
        else:
            server_address = (server_address, 42217)

    sock.connect(server_address)
except socket.error as e:
    sys.stderr.write(str(e) + ': ' + str(server_address) + '\n')
    sys.exit(1)

buffer = ''
max = -1
try:
    sock.settimeout(5)
    sock.sendall('STATS\n'.encode())
    while max == -1 or len(buffer.split('\n')) < max:
        buffer += sock.recv(1024).decode()
        if max == -1:
            # the first line contains the number of following lines
            max = int(buffer.split(' ')[0]) + 1
except socket.error as e:
    sys.stderr.write(str(e) + '\n')
    sys.exit(1)

sock.close()
print('<<<rrdcached>>>')
print(buffer.rstrip('\n'))