$search
00001 #! /usr/bin/env python 00002 00003 from twisted.internet.defer import inlineCallbacks, DeferredLock 00004 from twisted.internet import reactor 00005 import state_publisher 00006 import system 00007 00008 class IpRule: 00009 """Manage the ip rules for a given priority level. 00010 00011 There are two intended uses: 00012 00013 1. There is a single dynamic rule at that priority level. Use set to 00014 change it. 00015 00016 2. There are multiple static rules at that level. Use add to add each 00017 one of them, and flush to remove them all. 00018 """ 00019 def __init__(self, priority): 00020 self._lock = DeferredLock() 00021 self.priority = str(priority) 00022 self.is_shutdown = False 00023 self.state_pub = state_publisher.StatePublisher(()) 00024 self.flush() 00025 reactor.addSystemEventTrigger('before', 'shutdown', self._shutdown) 00026 00027 @inlineCallbacks 00028 def flush(self): 00029 self.prev = () # Won't hurt to do it here, just means we might not 00030 # delete during ongoing set. Ensures it is set at 00031 # end of construction. 00032 yield self._lock.acquire() 00033 try: 00034 retcode = None 00035 while not retcode: 00036 retcode = yield system.system(system.Quiet, 'ip', 'rule', 'del', 'priority', self.priority) 00037 finally: 00038 self._lock.release() 00039 00040 @inlineCallbacks 00041 def _set_or_add(self, remove_old, args): 00042 """Changes/adds the rule. Call with no arguments to clear the rule.""" 00043 if self.is_shutdown: 00044 return 00045 yield self._lock.acquire() 00046 try: 00047 # Are we replacing a rule with itself? 00048 if remove_old and self.prev and args == self.prev: 00049 return 00050 # Add new rule 00051 if args: 00052 yield system.system('ip', 'rule', 'add', "priority", self.priority, *args) 00053 # Add remove old rule 00054 if self.prev and remove_old: 00055 yield system.system('ip', 'rule', 'del', "priority", self.priority, *self.prev) 00056 self.prev = args 00057 self.state_pub.set(args) 00058 00059 finally: 00060 self._lock.release() 00061 00062 def set(self, *args): 00063 """Changes/adds the rule. Call with no arguments to clear the rule.""" 00064 self._set_or_add(True, args) 00065 00066 def add(self, *args): 00067 """Adds a rule without removing the existing one.""" 00068 self._set_or_add(False, args) 00069 00070 @inlineCallbacks 00071 def _shutdown(self): 00072 yield self.flush() 00073 self.is_shutdown = True