Coverage for pnr_observer.py : 78%

############################################################################# 

# 

# Manage the flash lists and site flags, alerting surrogates of updates 

# 

# Copyright (C) 2015-2016 Menlo Security, Inc. 

# All rights reserved. 

# 

############################################################################# 

import tornado 

import tornado.ioloop as ioloop 

from safly.config import config 

from safly.logger import SafelyLoggerMixin 

from safly.pnr_helper import PNRHelper 

class PnrObserver(SafelyLoggerMixin): 

def __init__(self, io_loop=None): 

super(PnrObserver, self).__init__('sm-pnr-observer', 

auto_prefix=True) 

self._io_loop = io_loop or ioloop.IOLoop.instance() 

self._pc = ioloop.PeriodicCallback(self._on_periodic, 

config.getint('surrogate_manager', 

'pnr_refresh_secs') * 1000) 

self._pc.start() 

self._observers_by_tid = {} 

# Cache of site flags 

self._site_flags = None 

# Cache of flash lists 

self._flash_lists = {} 

@staticmethod 

def instance(io_loop=None): 

if not hasattr(PnrObserver, '_instance'): 

PnrObserver._instance = PnrObserver(io_loop) 

return PnrObserver._instance 

@tornado.gen.coroutine 

def _on_periodic(self): 

new_site_flags = yield tornado.gen.Task(PNRHelper().get_request, 

'/api/v1/pe/site_flags') 

site_flags_changed = new_site_flags != self._site_flags 

if site_flags_changed: 

self._site_flags = new_site_flags 

self.log.info({'count': len(new_site_flags)}, 

event='site-flags-changed') 

for tid in self._observers_by_tid.keys(): 

flash_list = yield tornado.gen.Task( 

PNRHelper().post_request,'/api/v2/pe/flash_site_list', {'tid': tid}) 

# Check that the tid is still in use after the yield. 

if tid in self._observers_by_tid: 

flash_list_changed = self._flash_lists[tid] != flash_list 

if flash_list_changed: 

self._flash_lists[tid] = flash_list 

self.log.info({'count': len(flash_list), 'tid': tid}, 

event='flash-list-changed') 

if flash_list_changed or site_flags_changed: 

for observer in self._observers_by_tid[tid]: 

try: 

observer({'flash_list': self._flash_lists[tid], 

'site_flags': self._site_flags}) 

except Exception: 

self.log.exception({}, event='pnr-callback-exception') 

@tornado.gen.coroutine 

def register(self, tid, observer_callback): 

if tid in self._observers_by_tid: 

self._observers_by_tid[tid].add(observer_callback) 

raise tornado.gen.Return({'flash_list': self._flash_lists[tid], 

'site_flags': self._site_flags}) 

if self._site_flags is None: 

self._site_flags = yield tornado.gen.Task(PNRHelper().get_request, 

'/api/v1/pe/site_flags') 

self.log.info({'count': len(self._site_flags)}, 

event='site-flags-changed') 

self._flash_lists[tid] = yield tornado.gen.Task( 

PNRHelper().post_request, '/api/v2/pe/flash_site_list', {'tid': tid}) 

self.log.info({'count': len(self._flash_lists[tid]), 'tid': tid}, 

event='flash-list-changed') 

# Watch for race here - another surrogate could 

# have registered while we yielded 

self._observers_by_tid.setdefault(tid, set()).add(observer_callback) 

raise tornado.gen.Return({'flash_list': self._flash_lists[tid], 

'site_flags': self._site_flags}) 

def deregister(self, tid, func): 

if not tid in self._observers_by_tid: 

return 

if func in self._observers_by_tid[tid]: 

self._observers_by_tid[tid].remove(func) 

else: 

self.log.warn({'tid': tid}, event='deregister-no-entry') 

if not self._observers_by_tid[tid]: 

self._flash_lists.pop(tid) 

self._observers_by_tid.pop(tid)