def test_registration(self): class TestMixin: pass self.assertTrue(Mixin.register('mixin', TestMixin)) self.assertTrue(Mixin.unregister('mixin')) self.assertFalse(Mixin.unregister('mixin'))
def test_underunder_enter_and_underunder_exit(self): class A(object): pass class B(A): pass for mixin in self.expected_mixins: if mixin: mixee = Mixin(A(), mixin) mixee.__enter__() self.assertTrue(mixin in A.__bases__) self.assertFalse(mixin in B.__bases__) mixee.__exit__() self.assertFalse(mixin in A.__bases__) self.assertFalse(mixin in B.__bases__)
def test_get_mixin(self): keys = 'hi!', 'i have', 'heard of?', None actual_mixins = [Mixin.get_mixin(key) for key in keys] for ex, ac in zip(self.expected_mixins, actual_mixins): self.assertEqual(ex, ac)
self.default_validator() assert payload['type'] assert payload['type'] in valid_payload_types except AssertionError: raise RequestError("Invalid payload: %s" % payload) return fn(self, **kwargs) return validator class Query: """ Query handler """ @validate def execute(self, host=None): """ Handle an inbound query """ key, value = self.packet.payload['type'], self.packet.payload['value'] result = [] for host in self._r.keys(pattern="host_*"): host = json.JSONDecoder().decode(self._r.get(host)) query_values = host.get(key, None) if query_values: if value in query_values: result.append(host) return result Mixin.register("heard of?", Query)
""" Validates a packet""" payload = self.packet.payload try: valid_payload_types = ('hosts', 'services', 'addresses', 'disks') self.default_validator() print payload assert payload['type'] assert payload['type'] in valid_payload_types assert type(payload['export']) == list assert len(payload['export']) > 0 except AssertionError: raise RequestError("Invalid payload: %s" % payload) return fn(self, **kwargs) return validator class Export: """ Export handler """ @validate def execute(self, host=None): """ handle an inbound export """ if not host: host = self.packet.payload['host'] key = 'host_%s' % host return self._store(key, self.packet.payload) Mixin.register("i have", Export)
from pynpoint.errors import RequestError def validate(fn): """ Validation decorator """ def validator(self, **kwargs): """ Validates a packet""" payload = self.packet.payload try: self.default_validator() assert payload['host'] or kwargs['host'] except AssertionError: raise RequestError("Invalid payload: %s" % payload) return fn(self, **kwargs) return validator class Announcement: """ Announcement RequestHandler mixin """ @validate def execute(self, host=None): """ handle a hello """ if not host: host = self.packet.payload['host'] key = 'host_%s' % host return self._store(key, self.packet.payload) Mixin.register("hi!", Announcement)
def handle(self): mixin = Mixin.get_mixin(self.packet.request_type) with Mixin(self, mixin) as handler: # handler.execute is really self, with mixin mixed in. return handler.execute()