# coding=utf-8

"""
Collects data from squid servers

#### Dependencies

"""

import re
import socket
import diamond.collector


class SquidCollector(diamond.collector.Collector):

    def __init__(self, *args, **kwargs):
        super(SquidCollector, self).__init__(*args, **kwargs)
        self.squid_hosts = {}

        host_pattern = re.compile("(([^@]+)@)?([^:]+)(:([0-9]+))?")
        self.stat_pattern = re.compile("^([^ ]+) = ([0-9\.]+)$")

        for host in self.config['hosts']:
            matches = host_pattern.match(host)

            if matches.group(5):
                port = matches.group(5)
            else:
                port = 3128

            if matches.group(2):
                nick = matches.group(2)
            else:
                nick = port

            self.squid_hosts[nick] = {
                'host': matches.group(3),
                'port': int(port)
            }

    def get_default_config_help(self):
        config_help = super(SquidCollector, self).get_default_config_help()
        config_help.update({
            'hosts': 'List of hosts to collect from. Format is '
            + '[nickname@]host[:port], [nickname@]host[:port], etc',
        })
        return config_help

    def get_default_config(self):
        """
        Returns the default collector settings
        """
        config = super(SquidCollector, self).get_default_config()
        config.update({
            'hosts': ['localhost:3128'],
            'path': 'squid',
        })
        return config

    def _getData(self, host, port):
        try:
            squid_sock = socket.socket()
            squid_sock.connect((host, int(port)))
            squid_sock.settimeout(0.25)
            squid_sock.sendall("GET cache_object://localhost/counters HTTP/1.0"
                               + "\r\nHost: localhost\r\nAccept: */*\r\nConnec"
                               + "tion: close\r\n\r\n")

            fulldata = ''

            while True:
                data = squid_sock.recv(1024)
                if not data:
                    break
                fulldata = fulldata + data
        except Exception, e:
            self.log.error('Couldnt connect to squid: %s', e)
            return None
        squid_sock.close()

        return fulldata

    def collect(self):
        for nickname in self.squid_hosts.keys():
            squid_host = self.squid_hosts[nickname]

            fulldata = self._getData(squid_host['host'],
                                     squid_host['port']).splitlines()

            for data in fulldata:
                matches = self.stat_pattern.match(data)
                if matches:
                    self.publish_counter("%s.%s" % (nickname,
                                                    matches.group(1)),
                                         float(matches.group(2)))

@@ -2515,16 +2515,90 @@
 'port'%5D)
+%0A%0A            if fulldata is not None:%0A                fulldata = fulldata
 .splitli
@@ -2612,24 +2612,28 @@
             
+    
 for data in 
@@ -2638,24 +2638,28 @@
 n fulldata:%0A
+    
             
@@ -2702,16 +2702,20 @@
 h(data)%0A
+    
         
@@ -2750,24 +2750,28 @@
             
+    
 self.publish
@@ -2848,24 +2848,28 @@
             
+    
 matches.grou
@@ -2875,16 +2875,20 @@
 up(1)),%0A
+    
         

modify squid collector to correctly handle unresponsive squid instance