#!/usr/bin/env python                                                              

import re
import sys
import socket
import time
import pprint

sock_main = '/var/run/haproxy/admin.sock'

sock_prefix = '/var/run/haproxy/stats_'
sock_suffix = '.sock'

stat_strings = ['pxname', 'svname', 'status', 'check_status']

class HAPStats(object):

    info                = dict()    # general info storage                         
    column_names        = []        # column names                                 
    stats_aggregated    = dict()    # aggregated stats                             

    def _check_int(self, s):
        if len(s) > 0 and s[0] in ('-', '+'):
            return s[1:].isdigit()
        return s.isdigit()

    def get_info(self):
        self.sock_main = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock_main.settimeout(10)
        try:
            self.sock_main.connect(sock_main)
            self.sock_main.sendall('show info\n')
            data = ""
            while True:
                buf = self.sock_main.recv(16)
                if buf:
                    data += buf
                else:
                    break
            self.sock_main.close()
        except:
           data = ""
        finally:
           self.sock_main.close()

        # parse response line by line                                              
        for line in data.split('\n'):
            m = re.match('([a-zA-Z0-9_\\-]+):\s([a-zA-Z0-9_\s\B]+)', line)

            # if there was a match, update our list containing info data           
            if m != None:
                key = m.group(1).lower()
                value = m.group(2)

                # if value is number, convert to int type,                         
                # keep string otherwise                                            
                if self._check_int(value):
                    self.info[key] = int(value)

                else:
                    self.info[key] = value

        return self.info

    def __stats_read(self, socket_name):
        self.sock_main = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock_main.settimeout(10)
        try:
            self.sock_main.connect(socket_name)
            self.sock_main.sendall('show stat\n')
            response = ""
            while True:
                buf = self.sock_main.recv(16)
                if buf:
                    response += buf
                else:
                    break
            self.sock_main.close()
        except:
           response = ""
        finally:
           self.sock_main.close()
        return response

    def __stats_aggregate_init(self, header):

        m = re.match('#\spxname,svname.*', header)

        # check if provided header contains data from haproxy stats                
        if m is not None:

            for column_name in header[2:].split(','):

                if len(column_name) > 0:
                    self.column_names.append(column_name)

    def get_stats_aggregated(self,
            sock_prefix = sock_prefix,
            sock_suffix = sock_suffix
            ):

        # load general info from haproxy's main socket                             
        info = self.get_info()

        # parse header, create list with column names                              

        # iterate processes                                                        
        for proc_num in range(0, info['nbproc']):
            sock_name = sock_prefix + str(proc_num) + sock_suffix

            # initialize line counter                                              
            line_num = 0

            # parse statistics, line by line                                       
            for line in self.__stats_read(sock_name).split('\n'):

                if proc_num == 0 and line_num == 0:
                    self.__stats_aggregate_init(line)

                # skip comments                                                    
                if len(line) > 0 and line[0] == '#':
                    continue

                else:
                    items = line.split(',')

                    if len(items) > 0:

                        # initialize helper variables                              
                        item_index = 0
                        pxname = ''
                        svname = ''

                        for item in items:

                            if item_index == 0:
                                pxname = item
                                if pxname not in self.stats_aggregated:
                                    self.stats_aggregated[pxname] = dict()

                            elif item_index == 1:
                                svname = item
                                if svname not in self.stats_aggregated[pxname]:
                                    self.stats_aggregated[pxname][svname] = dict()

                            elif item_index > 1 and item_index < len(self.column_names):

                                column_name = self.column_names[item_index]

                                if column_name not in self.stats_aggregated[pxname][svname]:

                                    if self._check_int(item):
                                        self.stats_aggregated[pxname][svname][column_name] = int(item)
                                    else:
                                        self.stats_aggregated[pxname][svname][column_name] = item

                                else:
                                    if self._check_int(item):
                                        try:
                                            if self.stats_aggregated[pxname][svname][column_name] == "":
                                                self.stats_aggregated[pxname][svname][column_name] = 0

                                            self.stats_aggregated[pxname][svname][column_name] = self.stats_aggregated[pxname][svname][column_name] + int(item)

                                        except TypeError:
                                            continue

                                    else:
                                        self.stats_aggregated[pxname][svname][column_name] = item

                            item_index = item_index + 1


hstats = HAPStats()
info = hstats.get_info()
hstats.get_stats_aggregated()

for pxname in hstats.stats_aggregated:
    for svname in hstats.stats_aggregated[pxname]:
        sys.stdout.write(pxname + ',' + svname)
        for column in hstats.column_names[2:]:
            try:
                sys.stdout.write(',' + str(hstats.stats_aggregated[pxname][svname][column]))

            except KeyError:
                continue

        sys.stdout.write('\n')
