Major Refactoring

Author: pythonhacker

config.py

@@ -0,0 +1,268 @@
+import json
+import time
+import random
+import operator
+import sys
+import os
+from utils import enum
+
+# Configurations
+
+# Rotation Policies
+Policy = enum('ROTATION_RANDOM',
+              # Least recently used
+              'ROTATION_LRU',
+              # Switch to another region
+              'ROTATION_NEW_REGION',
+              # LRU + New region
+              'ROTATION_LRU_NEW_REGION')
+
+region_dict = {2: 'Dallas',
+               3: 'Fremont',
+               4: 'Atlanta',
+               6: 'Newark',
+               7: 'London',
+               8: 'Tokyo',
+               9: 'Singapore',
+               10: 'Frankfurt'}
+
+
+email_template = """
+
+I just switched a proxy node in the proxy infrastructure. Details are below.
+
+In: %(label)s, %(proxy_in)s
+Out: %(label)s, %(proxy_out)s
+
+Region: %(region)s
+
+-- Linode proxy daemon
+
+"""
+
+# Post process command
+post_process_cmd_template = """ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null %s@%s "%s" """
+iptables_restore_cmd = "sudo iptables-restore < /etc/iptables.rules"
+squid_restart_cmd = "sudo squid3 -f /etc/squid3/squid.conf"
+
+
+class ProxyConfig(object):
+    """ Class representing configuration of crawler proxy infrastructure """
+
+    def __init__(self, cfg='proxy.conf'):
+        """ Initialize proxy config from the config file """
+
+        self.parse_config(cfg)
+        # This is a file with each line of the form
+        # IPV4 address, datacenter code, linode-id, switch_in timestamp, switch_out timestamp
+        # E.g: 45.79.91.191, 3, 1446731065, 144673390
+        try:
+            proxies = map(lambda x: x.strip().split(','), open(self.proxylist).readlines())
+            # Proxy IP to (switch_in, switch_out) timestamp mappings
+            self.proxy_dict = {}
+            # Proxy IP to enabled mapping
+            self.proxy_state = {}
+            self.process_proxies(proxies)
+        except (OSError, IOError), e:
+            print e
+            sys.exit("Fatal error, proxy list input file " + self.proxylist + " not found!")
+        except ValueError, e:
+            print e
+            print self.proxylist + " is empty or has junk values"
+            
+        try:
+            self.proxy_template = open(self.lb_template).read()
+        except (OSError, IOError), e:
+            print e
+            sys.exit("Fatal error, template config input file " + template_file + " not found!")
+
+    def parse_config(self, cfg):
+        """ Parse the configuration file and load config """
+
+        self.config = json.load(open(cfg))
+        for key,value in self.config.items():
+            # Set attribute locally
+            setattr(self, key, value)
+
+        # Do some further processing
+        self.frequency = float(self.frequency)*3600.0
+        self.policy = eval('Policy.' + self.policy)
+        
+    def get_proxy_ips(self):
+        """ Return all proxy IP addresses as a list """
+
+        return self.proxy_state.keys()
+
+    def get_active_proxies(self):
+        """ Return a list of all active proxies as a list """
+
+        return map(self.proxy_dict.get, filter(self.proxy_state.get, self.proxy_state.keys()))
+        
+    def process_proxies(self, proxies):
+        """ Process the proxy information to create internal dictionaries """
+
+        # Prepare the proxy region dict
+        for proxy_ip, region, proxy_id, switch_in, switch_out in proxies:
+            # If switch_in ==0: put current time
+            if int(float(switch_in))==0:
+                switch_in = int(time.time())
+            if int(float(switch_out))==0:
+                switch_out = int(time.time())
+
+            self.proxy_dict[proxy_ip] = [proxy_ip, int(region), proxy_id, int(float(switch_in)), int(float(switch_out))]
+            self.proxy_state[proxy_ip] = True
+
+        print 'Processed',len(self.proxy_state),'proxies.'
+
+    def get_proxy_for_rotation(self,
+                               use_random=False,
+                               least_used=False,
+                               region_switch=False,
+                               input_region=3):
+        """ Return a proxy IP address for rotation using the given settings. The
+        returned proxy will be replaced with a new proxy.
+
+        @use_random - Means returns a random proxy from the current active list
+        @least_used - Returns a proxy IP which is the oldest switched out one
+        so we keep the switching more or less democratic.
+        @region_switch - Returns a proxy which belongs to a different region
+        from the new proxy.
+        @input_region - The region of the new proxy node - defaults to Fremont, CA.
+        
+        Note that if use_random is set to true, the other parameters are ignored.
+        
+        """
+
+        active_proxies = self.get_active_proxies()
+        print 'Active proxies =>',active_proxies
+        
+        if use_random:
+            # Pick a random proxy IP
+            proxy = random.choice(active_proxies)
+            print 'Returning proxy =>',proxy
+            proxy_ip = proxy[0]
+            
+            # Remove it from every data structure
+            self.switch_out_proxy(proxy_ip)
+            return proxy
+
+        if least_used:
+            # Pick the oldest switched out proxy i.e one
+            # with smallest switched out value
+            proxies_used = sorted(active_proxies,
+                                  key=operator.itemgetter(-1))
+
+            print 'Proxies used =>',proxies_used
+            
+            if region_switch:
+                # Find the one with a different region from input
+                for proxy, reg, pi, si, so in proxies_used:
+                    if reg != input_region:
+                        print 'Returning proxy',proxy,'from region',reg
+                        self.switch_out_proxy(proxy)                        
+                        return proxy
+
+            # If all regions are already in use, pick the last used
+            # proxy anyway
+            return proxies_used[0][0]
+            
+        if region_switch:
+            # Pick a random proxy not in the input region
+            proxies = active_proxies
+            random.shuffle(proxies)
+
+            for proxy, reg, pi, si, so in proxies:
+                if reg != input_region:
+                    print 'Returning proxy',proxy,'from region',reg
+                    self.switch_out_proxy(proxy)                    
+                    return proxy
+
+    def __getattr__(self, name):
+        """ Return from local, else written from config """
+
+        try:
+            return self.__dict__[name]
+        except KeyError:
+            return self.config.get(name)
+        
+    def switch_out_proxy(self, proxy):
+        """ Switch out a given proxy IP """
+
+        # Disable it
+        self.proxy_state[proxy] = False
+        # Mark its switched out timestamp
+        self.proxy_dict[proxy][-1] = int(time.time())
+
+    def switch_in_proxy(self, proxy, proxy_id, region):
+        """ Switch in a given proxy IP """
+
+        # Mark its switched out timestamp
+        self.proxy_dict[proxy] = [proxy, int(region), proxy_id, int(time.time()), int(time.time())]
+        # Enable it
+        self.proxy_state[proxy] = True     
+
+    def get_active_regions(self):
+        """ Return unique regions for which proxies are active """
+
+        regions = set()
+        for proxy,region,pi,si,so in self.proxy_dict.values():
+            if self.proxy_state[proxy]:
+                regions.add(region)
+
+        return list(regions)
+        
+    def write(self, disabled=False):
+        """ Write current state to an output file """
+
+        lines = []
+        for proxy, reg, pi, si, so in self.proxy_dict.values():
+            if disabled or self.proxy_state[proxy]:
+                lines.append('%s,%s,%s,%s,%s\n' % (proxy, str(reg), str(pi), str(int(si)), str(int(so))))
+
+        open(self.proxylist,'w').writelines(lines)
+
+    def write_lb_config(self, disabled=False, test=False):
+        """ Write current proxy configuration into the load balancer config """
+
+        lines, idx = [], 1
+        # Shuffle
+        items = self.proxy_dict.values()
+        for i in range(10):
+            random.shuffle(items)
+            
+        for proxy, reg, pi, si, so in items:
+            if self.proxy_state[proxy]:
+                lines.append('\tserver  squid%d %s:8321 check inter 10000 rise 2 fall 5' % (idx, proxy))
+                idx += 1
+
+        squid_config = "\n".join(lines)
+        content = self.proxy_template % locals()
+        # Write to temp file
+        tmpfile = '/tmp/.haproxy.cfg'
+        open(tmpfile,'w').write(content)
+
+        # If running in test mode, don't do this!
+        if not test:
+            # Run as sudo
+            cmd = 'sudo cp %s %s; rm -f %s' % (tmpfile, self.lb_config, tmpfile)
+            os.system(cmd)
+
+        self.reload_lb()
+        return True
+
+    def reload_lb(self):
+        """ Reload the HAProxy load balancer """
+
+        return (os.system(self.lb_restart) == 0)
+
+    def get_proxy_id(self, proxy):
+        """ Given proxy return its id """
+
+        return self.proxy_dict[proxy][2]
+
+    def get_email_config(self):
+        """ Return email configuration """
+
+        return self.config['email']
+    
+