#!/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.connect(sock_main)                                          
        self.sock_main.sendall('show info\n')                                      
        data = self.sock_main.recv(4096)                                           
                                                                                   
        # 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.connect(socket_name)                                        
        self.sock_main.sendall('show stat\n')                                      
        response = self.sock_main.recv(4096)                                       
        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):                   
                                        self.stats_aggregated[pxname][svname][column_name] = self.stats_aggregated[pxname][svname][column_name] + int(item)
                                                                                
                                    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:]:                                  
            sys.stdout.write(',' + str(hstats.stats_aggregated[pxname][svname][column]))
                                                                                
        sys.stdout.write('\n')    
