Skip to content

Commit 7df8af4

Browse files
pythonhackerpythonhacker
committed
Major Refactoring
1 parent 55506af commit 7df8af4

File tree

1 file changed

+281
-0
lines changed

1 file changed

+281
-0
lines changed

‎base.py‎

Lines changed: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,281 @@
1+
import os
2+
import threading
3+
import signal
4+
import time
5+
import collections
6+
7+
import email_report
8+
from config import *
9+
10+
from utils import daemonize
11+
12+
class ProxyRotator(object):
13+
""" Proxy rotation, provisioning & re-configuration base class """
14+
15+
def __init__(self, cfg='proxy.conf', test_mode=False, rotate=False, region=None):
16+
self.config = ProxyConfig(cfg=cfg)
17+
print 'Frequency set to',self.config.frequency,'seconds.'
18+
# Test mode ?
19+
self.test_mode = test_mode
20+
# Event object
21+
self.alarm = threading.Event()
22+
# Clear the event
23+
self.alarm.clear()
24+
# Heartbeat file
25+
self.hbf = '.heartbeat'
26+
# Actual command used for doing stuff
27+
self.vps_command = None
28+
# If rotate is set, rotate before going to sleep
29+
if rotate:
30+
print 'Rotating a node'
31+
self.rotate(region=region)
32+
33+
signal.signal(signal.SIGTERM, self.sighandler)
34+
signal.signal(signal.SIGUSR1, self.sighandler)
35+
36+
def pick_region(self):
37+
""" Pick the region for the new node """
38+
39+
# Try and pick a region not present in the
40+
# current list of nodes
41+
regions = self.config.get_active_regions()
42+
# Shuffle current regions
43+
random.shuffle(self.config.region_ids)
44+
45+
for reg in self.config.region_ids:
46+
if reg not in regions:
47+
return reg
48+
49+
# All regions already present ? Pick a random one.
50+
return random.choice(self.config.region_ids)
51+
52+
def rotate(self, region=None):
53+
""" Rotate the configuration to a new node """
54+
55+
proxy_out_label = None
56+
# Pick the data-center
57+
if region == None:
58+
print 'Picking a region ...'
59+
region = self.pick_region()
60+
else:
61+
print 'Using supplied region',region,'...'
62+
63+
# Switch in the new linode from this region
64+
new_proxy, proxy_id = self.make_new_instance(region)
65+
66+
# Rotate another node
67+
if self.config.policy == Policy.ROTATION_RANDOM:
68+
proxy_out = self.config.get_proxy_for_rotation(use_random=True, input_region=region)
69+
elif self.config.policy == Policy.ROTATION_NEW_REGION:
70+
proxy_out = self.config.get_proxy_for_rotation(region_switch=True, input_region=region)
71+
elif self.config.policy == Policy.ROTATION_LRU:
72+
proxy_out = self.config.get_proxy_for_rotation(least_used=True, input_region=region)
73+
elif self.config.policy == Policy.ROTATION_LRU_NEW_REGION:
74+
proxy_out = self.config.get_proxy_for_rotation(least_used=True, region_switch=True,
75+
input_region=region)
76+
77+
# Switch in the new proxy
78+
self.config.switch_in_proxy(new_proxy, proxy_id, region)
79+
print 'Switched in new proxy',new_proxy
80+
# Write configuration
81+
self.config.write()
82+
print 'Wrote new configuration.'
83+
# Write new HAProxy LB template and reload ha proxy
84+
ret1 = self.config.write_lb_config()
85+
ret2 = self.config.reload_lb()
86+
87+
if ret1 and ret2:
88+
if proxy_out != None:
89+
print 'Switched out proxy',proxy_out
90+
proxy_out_id = int(self.config.get_proxy_id(proxy_out))
91+
92+
if proxy_out_id != 0:
93+
proxy_out_label = self.get_instance_label(proxy_out_id)
94+
print 'Removing switched out instance',proxy_out_id
95+
self.delete_instance(proxy_out_id)
96+
else:
97+
'Proxy id is 0, not removing proxy',proxy_out
98+
else:
99+
print 'Error - Did not switch out proxy as there was a problem in writing/restarting LB'
100+
101+
if proxy_out_label != None:
102+
# Get its label and assign it to the new linode
103+
print 'Assigning label',proxy_out_label,'to new instance',proxy_id
104+
time.sleep(5)
105+
self.update_instance(proxy_id,
106+
proxy_out_label,
107+
self.config.group)
108+
109+
# Post process the host
110+
print 'Post-processing',new_proxy,'...'
111+
self.post_process(new_proxy)
112+
self.send_email(proxy_out, proxy_out_label, new_proxy, region)
113+
114+
def post_process(self, ip):
115+
""" Post-process a switched-in host """
116+
117+
# Sleep a bit before sshing
118+
time.sleep(5)
119+
cmd = post_process_cmd_template % (self.config.user, ip, iptables_restore_cmd)
120+
print 'SSH command 1=>',cmd
121+
os.system(cmd)
122+
cmd = post_process_cmd_template % (self.config.user, ip, squid_restart_cmd)
123+
print 'SSH command 2=>',cmd
124+
os.system(cmd)
125+
126+
def provision(self, count=8, add=False):
127+
""" Provision an entirely fresh set of linodes after dropping current set """
128+
129+
if not add:
130+
self.drop()
131+
132+
num, idx = 0, 0
133+
134+
# If we are adding Linodes without dropping, start from current count
135+
if add:
136+
start = len(self.config.get_active_proxies())
137+
else:
138+
start = 0
139+
140+
for i in range(start, start + count):
141+
142+
# Do a round-robin on regions
143+
region = self.config.region_ids[idx % len(self.config.region_ids) ]
144+
try:
145+
ip, lid = self.make_new_instance(region)
146+
new_label = self.config.proxy_prefix + str(i+1)
147+
self.update_instance(int(lid),
148+
new_label,
149+
self.config.group)
150+
151+
num += 1
152+
except Exception, e:
153+
print 'Error creating instance',e
154+
155+
idx += 1
156+
157+
print 'Provisioned',num,' proxies.'
158+
# Save latest proxy information
159+
self.write_proxies()
160+
161+
def write_proxies(self):
162+
""" Write proxies to a file """
163+
164+
proxies_list = self.vps_command.get_proxies()
165+
# Randomize it
166+
for i in range(5):
167+
random.shuffle(proxies_list)
168+
169+
filename = self.config.proxylist
170+
print >> open(filename, 'w'), '\n'.join(proxies_list)
171+
print 'Saved current proxy configuration to {}'.format(filename)
172+
173+
def test(self):
174+
""" Function to be called in loop for testing """
175+
176+
proxy_out_label = ''
177+
region = self.pick_region()
178+
print 'Rotating proxy to new region',region,'...'
179+
# Make a test IP
180+
new_proxy, proxy_id = self.make_new_linode(region, test=True)
181+
proxy_out = self.config.get_proxy_for_rotation(least_used=True, region_switch=True,
182+
input_region=region)
183+
184+
if proxy_out != None:
185+
print 'Switched out proxy',proxy_out
186+
proxy_out_id = int(self.config.get_proxy_id(proxy_out))
187+
proxy_out_label = self.linode_cmd.get_label(proxy_out_id)
188+
189+
# Switch in the new proxy
190+
self.config.switch_in_proxy(new_proxy, proxy_id, region)
191+
print 'Switched in new proxy',new_proxy
192+
# Write new HAProxy LB template and reload ha proxy
193+
self.config.write_lb_config(test=True)
194+
self.send_email(proxy_out, proxy_out_label, new_proxy, region)
195+
196+
def stop(self):
197+
""" Stop the rotator process """
198+
199+
try:
200+
os.remove(self.hbf)
201+
# Signal the event
202+
self.alarm.set()
203+
return True
204+
except (IOError, OSError), e:
205+
pass
206+
207+
return False
208+
209+
def sighandler(self, signum, stack):
210+
""" Signal handler """
211+
212+
# This will be called when you want to stop the daemon
213+
self.stop()
214+
215+
def run(self):
216+
""" Run as a background process, rotating proxies """
217+
218+
# Touch heartbeat file
219+
open(self.hbf,'w').write('')
220+
# Fork
221+
print 'Daemonizing...'
222+
daemonize('rotator.pid',logfile='rotator.log', drop=True)
223+
print 'Proxy rotate daemon started.'
224+
count = 1
225+
226+
while True:
227+
# Wait on event object till woken up
228+
self.alarm.wait(self.config.frequency)
229+
status = self.alive()
230+
if not status:
231+
print 'Daemon signalled to exit. Quitting ...'
232+
break
233+
234+
print 'Rotating proxy node, round #%d ...' % count
235+
if self.test_mode:
236+
self.test()
237+
else:
238+
self.rotate()
239+
count += 1
240+
241+
sys.exit(0)
242+
243+
def create(self, region=3):
244+
""" Create a new instance for testing """
245+
246+
print 'Creating new instance in region',region,'...'
247+
new_proxy = self.make_new_instance(region, verbose=True)
248+
249+
return new_proxy
250+
251+
def send_email(self, proxy_out, label, proxy_in, region):
252+
""" Send email upon switching of a proxy """
253+
254+
print 'Sending email...'
255+
region = region_dict[region]
256+
content = email_template % locals()
257+
email_config = self.config.get_email_config()
258+
259+
email_report.email_report(email_config, "%s", content)
260+
261+
def alive(self):
262+
""" Return whether I should be alive """
263+
264+
return os.path.isfile(self.hbf)
265+
266+
def get_instance_label(self, instance_id):
267+
""" Return instance label given instance id """
268+
pass
269+
270+
def update_instance(self, instance_id, label, group=None):
271+
""" Update the meta-data for the instance """
272+
pass
273+
274+
def delete_instance(self, instance_id):
275+
""" Delete a given instance given its id """
276+
pass
277+
278+
def drop(self):
279+
""" Drop all instances in current configuration (except the LB) """
280+
pass
281+

0 commit comments

Comments
 (0)