class TASK(VTask):
LOOPLESS = True
OPT_PREFIX = 'my_task'
basicopt = option(default="spam")
opt_uscore = option(default="eggs")
opt_uscore2 = option(name="opt_uscore2", default="ham")
class Ros3DdevControllerService(VService):
"""Ros3D device controller services wrapper"""
system_config_file = option(default=SystemConfigLoader.DEFAULT_PATH,
help='Ros3D system configuration file')
config_file = option(default=ControllerConfigLoader.DEFAULT_PATH,
help='Ros3D controller configuration file')
def __init__(self, *args, **kwargs):
super(Ros3DdevControllerService, self).__init__(*args, **kwargs)
self.logger.debug('loading controller configuration from %s',
self.options.config_file)
self.config = ControllerConfigLoader(self.options.config_file)
self.logger.debug('loading system configuration from %s',
self.options.system_config_file)
self.system_config = SystemConfigLoader(self.options.system_config_file)
self.controller = Controller()
self.controller.set_snapshots_location(self.config.get_snapshots_location())
def initLogging(self):
"""Setup logging to stderr"""
logging.basicConfig(level=self.loglevel, stream=sys.stderr,
format='%(thread)d %(asctime)s - %(name)s - %(levelname)s - %(message)s')
captureWarnings(True)
class CPPServerTask(ThriftServerTask):
"""Spin up a thrift TNonblockingServer in a sparts worker thread"""
# DEFAULT_HOST = '0.0.0.0'
DEFAULT_PORT = 0
OPT_PREFIX = 'thrift'
bound_host = bound_port = None
# host = option(default=lambda cls: cls.DEFAULT_HOST, metavar='HOST',
# help='Address to bind server to [%(default)s]')
port = option(default=lambda cls: cls.DEFAULT_PORT,
type=int,
metavar='PORT',
help='Port to run server on [%(default)s]')
num_threads = option(name='threads',
default=4,
type=int,
metavar='N',
help='Server Worker Threads [%(default)s]')
def initTask(self):
"""Overridden to bind sockets, etc"""
super(CPPServerTask, self).initTask()
self._stopped = False
self.server = TCppServer(self.processor)
# TODO: There is no API to set host yet.
# self.server.setHost(self.host)
self.server.setPort(self.port)
self.server.setNWorkerThreads(self.num_threads)
self.server.setup()
self.bound_addrs = [self.server.getAddress()]
for addrinfo in self.bound_addrs:
self.bound_host, self.bound_port = addrinfo[0:2]
self.logger.info(
"%s Server Started on %s", self.name,
self._fmt_hostport(self.bound_host, self.bound_port))
def _fmt_hostport(self, host, port):
if ':' in host:
return '[%s]:%d' % (host, port)
else:
return '%s:%d' % (host, port)
def stop(self):
"""Overridden to tell the thrift server to shutdown asynchronously"""
self.server.stop()
def _runloop(self):
"""Overridden to execute TNonblockingServer's main loop"""
self.server.serve()
class DBusServiceTask(DBusTask):
"""Glue Task for exporting this VService over dbus"""
OPT_PREFIX = 'dbus'
BUS_NAME = None
BUS_CLASS = VServiceDBusObject
bus_name = option(default=lambda cls: cls.BUS_NAME,
metavar='NAME',
help='Bus Name. Should be something like '
'"com.sparts.AwesomeService"')
replace = option(action='store_true',
type=bool,
default=False,
help='Replace, and enable replacing of this service')
queue = option(action='store_true',
type=bool,
default=False,
help='If not --{task}-replace, will wait to take '
'this bus name')
dbus_service = None
def initTask(self):
super(DBusServiceTask, self).initTask()
assert self.bus_name is not None, \
"You must pass a --{task}-bus-name"
def start(self):
self.bus = dbus.SessionBus(private=True)
self.dbus_service = dbus.service.BusName(self.bus_name, self.bus,
self.replace, self.replace,
self.queue)
self.addHandlers()
super(DBusServiceTask, self).start()
def addHandlers(self):
self.sparts_dbus = self.BUS_CLASS(self)
task = self.service.getTask(FB303HandlerTask)
if task is not None:
self.fb303_dbus = FB303DbusService(self.dbus_service, task,
self.service.name)
def stop(self):
if self.dbus_service is not None:
self.dbus_service = None
#self.bus.close()
super(DBusServiceTask, self).stop()
class HostCheckTask(PeriodicTask):
INTERVAL = 5
check_name = option(default=socket.getfqdn(),
type=str,
help='Name to check [%(default)s]')
def execute(self, *args, **kwargs):
self.logger.info("LOOKUP %s => %s", self.check_name,
socket.gethostbyname(self.check_name))
class PeriodicTask(VTask):
"""Task that executes `execute` at a specified interval
You must either override the `INTERVAL` (seconds) class attribute, or
pass a --{OPT_PREFIX}-interval in order for your task to run.
"""
INTERVAL = None
execute_duration_ms = samples(
windows=[60, 240],
types=[SampleType.AVG, SampleType.MAX, SampleType.MIN])
n_iterations = counter()
n_slow_iterations = counter()
n_try_later = counter()
interval = option(type=float,
metavar='SECONDS',
default=lambda cls: cls.INTERVAL,
help='How often this task should run [%(default)s] (s)')
def execute(self, context=None):
"""Override this to perform some custom action periodically."""
self.logger.debug('execute')
def initTask(self):
super(PeriodicTask, self).initTask()
assert self.interval is not None
# Register an event that we can more smartly wait on in case shutdown
# is requested while we would be `sleep()`ing
self.stop_event = Event()
def stop(self):
self.stop_event.set()
super(PeriodicTask, self).stop()
def _runloop(self):
t0 = time.time()
while not self.service._stop:
try:
self.execute()
except TryLater:
self.n_try_later.increment()
continue
self.n_iterations.increment()
self.execute_duration_ms.add((time.time() - t0) * 1000)
to_sleep = (t0 + self.interval) - time.time()
if to_sleep > 0:
if self.stop_event.wait(to_sleep):
return
else:
self.n_slow_iterations.increment()
t0 = time.time()
class NBServerTask(ThriftServerTask):
"""Spin up a thrift TNonblockingServer in a sparts worker thread"""
DEFAULT_HOST = '0.0.0.0'
DEFAULT_PORT = 0
OPT_PREFIX = 'thrift'
bound_host = bound_port = None
host = option(default=lambda cls: cls.DEFAULT_HOST, metavar='HOST',
help='Address to bind server to [%(default)s]')
port = option(default=lambda cls: cls.DEFAULT_PORT,
type=int, metavar='PORT',
help='Port to run server on [%(default)s]')
num_threads = option(name='threads', default=10, type=int, metavar='N',
help='Server Worker Threads [%(default)s]')
def initTask(self):
"""Overridden to bind sockets, etc"""
super(NBServerTask, self).initTask()
self._stopped = False
self.socket = TServerSocket(self.host, self.port)
self.server = TNonblockingServer(self.processor, self.socket,
threads=self.num_threads)
self.server.prepare()
self.bound_host, self.bound_port = \
self.server.socket.handle.getsockname()
self.logger.info("%s Server Started on %s:%s",
self.name, self.bound_host, self.bound_port)
def stop(self):
"""Overridden to tell the thrift server to shutdown asynchronously"""
self.server.stop()
self.server.close()
self._stopped = True
def _runloop(self):
"""Overridden to execute TNonblockingServer's main loop"""
while not self.server._stop:
self.server.serve()
while not self._stopped:
time.sleep(0.1)
class TsdDBQueueTask(QueueTask):
OPT_PREFIX = 'sqlite'
workers = 1
# DEFAULT_SQLITE_FILE = '/var/lib/atcd.db'
DEFAULT_SQLITE_FILE = '/home/pi/traffic-shark-service/tsd.db'
sqlite_file = option(
default=DEFAULT_SQLITE_FILE,
metavar='SQLITE_FILE',
help='Location to store the sqlite3 db [%(default)s]',
name='file',
)
def initTask(self):
super(TsdDBQueueTask, self).initTask()
try:
self.sqlite_manager = SQLiteManager(self.sqlite_file, self.logger)
except OperationalError:
self.logger.exception(
'Unable to initialize DB from file "{0}"'.format(
self.sqlite_file))
raise
def execute(self, item, context):
try:
obj, action = item
except ValueError:
self.logger.exception('Error executing on item: {0}'.format(item))
return
try:
func = getattr(self.sqlite_manager, action)
except AttributeError:
self.logger.exception(
'unable to run action, {0}, no such method'.format(action))
raise
try:
if isinstance(obj, tuple):
func(*obj)
else:
func(obj)
except OperationalError:
self.logger.exception("Unsupported operation")
return
def get_saved_profiles(self):
return self.sqlite_manager.get_saved_profiles()
def get_saved_mcontrols(self):
return self.sqlite_manager.get_saved_mcontrols()
class Reconnector(PeriodicTask, DBusTask, NMDBusHelper):
OPT_PREFIX = 'reconnect'
INTERVAL = 5.0
LOOPLESS = False
ssid = option(required=True,
type=str,
metavar='SSID',
help='SSID to force reconnection with')
def initTask(self):
super(Reconnector, self).initTask()
self.sbus = dbus.SystemBus(mainloop=self.mainloop_task.dbus_loop)
def execute(self):
# Find the connection settings for the SSID
wanted_conn = self.get_wireless_conn_by_ssid(self.ssid)
for device in self.iter_wireless_devices():
# If the active access point is already correct, we're done.
try:
active_ap = self.get_device_active_ap(device)
active_ap_ssid = self.get_ap_ssid(active_ap)
except dbus.DBusException:
self.logger.exception("Error getting active AP")
raise TryLater("Error getting active AP", after=2.5)
self.logger.debug('Active AP SSID is %s', active_ap_ssid)
if active_ap_ssid == self.ssid:
break
# Find the actual AP for the wanted SSID
wanted_ap = self.get_device_ap_by_ssid(device, self.ssid)
# Connect the device, with saved connection settings, to the AP
t0 = time.time()
self.logger.info("Connecting to %s...", self.ssid)
nm = self.get_networkmanager()
active_connection = nm.ActivateConnection(
wanted_conn.object_path,
device.object_path,
wanted_ap.object_path,
)
self.logger.info("Connection took %.2fs", time.time() - t0)
self.logger.info("result was %s", active_connection)
class MQTTTask(TornadoTask):
OPT_PREFIX = 'mqtt'
port = option(default=1883, help='Broker port')
host = option(default='localhost', help='Broker host address')
topic = option(default='/parameters', help='Parameters topic')
def __init__(self, *args, **kwargs):
super(MQTTTask, self).__init__(*args, **kwargs)
self.adapter = None
self.client = mqtt.Client()
self.client.on_connect = self._on_connect
self.client.on_disconnect = self._on_disconnect
def start(self):
_log.debug('start')
self.ioloop.add_callback(self._try_connect)
ParametersStore.change_listeners.add(self.param_changed)
def stop(self):
_log.debug('stop')
if self.adapter is not None:
self.adapter.stop()
#self.client.disconnect()
ParametersStore.change_listeners.remove(self.param_changed)
def _try_connect(self):
"""Attempt connection, if not schedule a reconnect after timeout """
_log.debug('try connecting to broker')
try:
self._connect()
except socket.error:
_log.warning('failed to connect')
self._schedule_reconnect()
def _connect(self):
_log.debug('connect to %s:%d', self.host, int(self.port))
self.client.connect(self.host, self.port)
self.adapter = MQTTornadoAdapter(self.client)
def _schedule_reconnect(self):
"""Schedule reconnection attempt"""
# reconnect after 10s
timeout = 10
_log.debug('reconnect after %d seconds', timeout)
loop = self.ioloop
_log.debug('reconnect at %d (now %d)',
loop.time() + timeout, loop.time())
loop.add_timeout(loop.time() + timeout, self._try_connect)
def _on_connect(self, client, userdata, flags, rc):
"""CONNACK received"""
_log.debug('connected: %s', mqtt.connack_string(rc))
_log.debug('flags: %s', flags)
self.adapter.poll_writes()
def _on_disconnect(self, client, userdata, rc):
"""DISCONNECT received"""
_log.debug('disconnected: %s', mqtt.connack_string(rc))
self.adapter.stop()
self.adapter = None
self._schedule_reconnect()
def param_changed(self, param):
_log.debug('param_changed %s', param)
self.ioloop.add_callback(self._publish_param, param)
def _publish_param(self, param):
as_json = ParameterCodec(as_set=True).encode(param)
_log.debug('publish to %s: %s', self.topic, as_json)
self.client.publish(self.topic, as_json)
class QueueTask(VTask):
"""Task that calls `execute` for all work put on its `queue`"""
MAX_ITEMS = 0
WORKERS = 1
max_items = option(type=int,
default=lambda cls: cls.MAX_ITEMS,
help='Set a bounded queue length. This may '
'cause unexpected deadlocks. [%(default)s]')
workers = option(type=int,
default=lambda cls: cls.WORKERS,
help='Number of threads to spawn to work on items from '
'its queue. [%(default)s]')
execute_duration_ms = samples(
windows=[60, 240],
types=[SampleType.AVG, SampleType.MAX, SampleType.MIN])
n_trylater = counter()
n_completed = counter()
n_unhandled = counter()
def execute(self, item, context):
"""Implement this in your QueueTask subclasses"""
raise NotImplementedError()
def _makeQueue(self):
"""Override this if you need a custom Queue implementation"""
return queue.Queue(maxsize=self.max_items)
def initTask(self):
super(QueueTask, self).initTask()
self.queue = self._makeQueue()
self.counters['queue_depth'] = \
CallbackCounter(lambda: self.queue.qsize())
self._shutdown_sentinel = object()
def stop(self):
super(QueueTask, self).stop()
self.queue.put(self._shutdown_sentinel)
def submit(self, item):
"""Enqueue `item` into this task's Queue. Returns a `Future`"""
future = Future()
work = ExecuteContext(item=item, future=future)
self.queue.put(work)
return future
def map(self, items, timeout=None):
"""Enqueues `items` into the queue"""
futures = map(self.submit, items)
return [f.result(timeout) for f in futures]
def _runloop(self):
while not self.service._stop:
try:
item = self.queue.get(timeout=1.0)
if item is self._shutdown_sentinel:
self.queue.put(item)
break
except queue.Empty:
continue
# Create an ExecuteContext if we didn't have one
if isinstance(item, ExecuteContext):
context = item
item = context.item
context.raw_wrapped = False
else:
context = ExecuteContext(item=item)
context.raw_wrapped = True
try:
context.start()
result = self.execute(item, context)
self.work_success(context, result)
except TryLater:
self.work_retry(context)
except Exception as ex:
self.work_fail(context, ex)
finally:
self.queue.task_done()
def work_success(self, context, result):
self.n_completed.increment()
self.execute_duration_ms.add(context.elapsed * 1000.0)
context.set_result(result)
self.work_done(context)
def work_retry(self, context):
self.n_trylater.increment()
context.attempt += 1
self.work_done(context)
self.queue.put(context)
def work_fail(self, context, exception):
self.n_unhandled.increment()
self.execute_duration_ms.add(context.elapsed * 1000.0)
handled = context.set_exception(exception)
self.work_done(context)
if not handled:
raise
def work_done(self, context):
pass
class TsdThriftHandlerTask(ThriftHandlerTask):
DEFAULT_LAN = 'wlan0'
DEFAULT_WAN = 'eth0'
DEFAULT_IPTABLES = '/sbin/iptables'
DEFAULT_BURST_SIZE = 12000
DEFAULT_PCAP_PATH = '/tmp/pcaps/'
MODULE = TrafficSharkService
DEPS = [TsdDBQueueTask, TsdScapyTask]
OPT_PREFIX = 'tsd'
lan_name = option(
default=DEFAULT_LAN,
metavar='LAN',
help='name of the LAN interface [%(default)s]',
name='lan',
)
wan_name = option(
default=DEFAULT_WAN,
metavar='WAN',
help='name of the WAN interface [%(default)s]',
name='wan',
)
iptables = option(default=DEFAULT_IPTABLES,
metavar='IPTABLES',
help='location of the iptables binary [%(default)s]')
burst_size = option(
default=DEFAULT_BURST_SIZE,
metavar='BURST_SIZE',
type=int,
help='Amount of bytes that can be burst at a capped speed '
'[%(default)s]')
dont_drop_packets = option(
action='store_true',
help='[EXPERIMENTAL] Do not drop packets when going above max allowed'
' rate. Packets will be queued instead. Please mind that this'
' option will likely disappear in the future and is only provided'
' as a workaround until better longer term solution is found.',
)
@staticmethod
def factory():
#load protocols
regex_filter = re.compile('^[a-z|A-Z|_|-]+\.py$')
files = os.listdir('tsd/protocols')
for file in files:
if regex_filter.match(file) is not None:
try:
__import__('tsd.protocols.' + file[:-3])
except ImportError:
raise
# load backends
os_name = os.uname()[0]
klass = 'Tsd{0}Shaper'.format(os_name)
try:
if klass not in globals():
from_module_import_class(
'tsd.backends.{0}'.format(os_name.lower()), klass)
except AttributeError:
raise NotImplementedError('{0} is not implemented!'.format(klass))
except ImportError:
raise NotImplementedError('{0} backend is not implemented!'.format(
os_name.lower()))
return globals()[klass]
def initTask(self):
super(TsdThriftHandlerTask, self).initTask()
print "[start initTask]"
self.db_task = self.service.tasks.TsdDBQueueTask
self.scapy_task = self.service.tasks.TsdScapyTask
self.scapy_task.setupIface(self.lan_name)
self.lan = {'name': self.lan_name}
self.wan = {'name': self.wan_name}
self._links_lookup()
self.initialize_id_manager()
self.initialize_shaping_system()
# Map of MAC address to tc object that is currently
# being used to shape traffic from that device.
# {mac: {'ip': ip, 'tc': tc, 'is_shaping': True/False, online: True/False, 'last_time': timestamp}}
# {'AB:CD:EF:GH:9O': {'ip': "172.1.44.3", 'tc': TrafficControl(...), 'is_shaping': True, online: True, 'last_time': 1234567890123}}
self._machine_controls = {}
self._machine_shapings = {}
# Map of profile that is saved
# {id: {'name': name, 'tc_setting': tc_setting}, ...}
# {1: {'name': 'test', 'tc_setting': TrafficControlSetting(...)}, ...}
self._profiles = {}
self.logger.info('Restoring shaped & profile connection from DB')
self._restore_saved_profiles()
self._restore_saved_mcontrols()
def initialize_id_manager(self):
"""Initialize the Id Manager. This is architecture dependant as the
shaping subsystems may have different requirements.
"""
self.idmanager = IdManager(first_id=type(self).ID_MANAGER_ID_MIN,
max_id=type(self).ID_MANAGER_ID_MAX)
def initialize_shaping_system(self):
"""Initialize the shaping subsystem.
Each shaping platform should implement its own.
"""
raise NotImplementedError('Subclass should implement this')
def _restore_saved_profiles(self):
"""Restore the profiles from the sqlite3 db.
"""
# Isolate the things we are using eval on to reduce possible clownyness
# later on, also this way we don't have unused imports from importing
# blindly for eval
names = [
'Shaping', 'Profile', 'TrafficControlSetting', 'Loss', 'Delay',
'Corruption', 'Reorder'
]
globals = {
name: getattr(traffic_shark_thrift.ttypes, name)
for name in names
}
result = []
try:
results = self.db_task.get_saved_profiles()
except OperationalError:
self.logger.exception('Unable to perform DB operation')
# self.logger.info("profiles:{0}".format(results))
for result in results:
profile_name = result['name']
profile_tc_setting = eval(result['tc_setting'], globals)
self._profiles[profile_name] = profile_tc_setting
def _restore_saved_mcontrols(self):
"""Restore the shapings from the sqlite3 db.
"""
# Isolate the things we are using eval on to reduce possible clownyness
# later on, also this way we don't have unused imports from importing
# blindly for eval
names = [
'Shaping', 'TrafficControlSetting', 'Loss', 'Delay', 'Corruption',
'Reorder'
]
globals = {
name: getattr(traffic_shark_thrift.ttypes, name)
for name in names
}
result = []
try:
results = self.db_task.get_saved_mcontrols()
except OperationalError:
self.logger.exception('Unable to perform DB operation')
for result in results:
mac = result['mac']
self._machine_controls[mac] = {
'mac': result['mac'],
'ip': result['ip'],
'profile_name': result['profile_name'],
'is_capturing': result['is_capturing'],
'is_shaping': result['is_shaping'],
'online': result['online'],
'capture_filter': result['capture_filter'],
'last_update_time': result['last_time']
}
for mac in self._machine_controls:
mc = self._machine_controls[mac]
if mc['is_shaping']:
self.shapeMachine(mac)
if mc['is_capturing']:
self.startCapture(mac, mc.get("capture_filter"))
def run_cmd(self, cmd):
self.logger.info("Running {}".format(cmd))
return subprocess.call(shlex.split(cmd))
# Profiles
def _profile_instance(self, name, settings):
return Profile(
name=name,
tc_setting=settings,
)
def getProfiles(self):
self.logger.info('Request getProfiles')
# print self._profiles
return [
self._profile_instance(name, settings)
for name, settings in self._profiles.items()
]
def addProfile(self, profile):
self.logger.info('Request addProfile for name {0}'.format(
profile.name))
self.db_task.queue.put(
((profile.name, profile.tc_setting), 'add_profile'))
self._profiles[profile.name] = profile.tc_setting
mac_list = []
for mac in self._machine_shapings:
mshaping = self._machine_shapings[mac]
if mshaping['profile_name'] == profile.name:
mac_list.append(mac)
for mac in mac_list:
self.shapeMachine(mac) # reshape
return TrafficControlRc(code=ReturnCode.OK)
def removeProfile(self, name):
self.logger.info('Request removeProfile for name {0}'.format(name))
self.db_task.queue.put(((name), 'remove_profile'))
del self._profiles[name]
mac_list = []
for mac in self._machine_shapings:
mshaping = self._machine_shapings[mac]
if mshaping['profile_name'] == name:
mac_list.append(mac)
for mac in mac_list:
self.unshapeMachine(mac) # unshape
for mac in self._machine_controls:
mc = self._machine_controls[mac]
if mc.get('profile_name') and mc['profile_name'] == name:
del mc['profile_name']
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
# Machine Controls
def _scanAddress(self):
addr_str = os.popen(
"arp -a | grep -v incomplete | grep wlan0 | awk '{print \"{\\\"ip\\\":\\\"\"substr($2, 2, length($2) - 2)\"\\\",\\\"mac\\\":\\\"\"$4\"\\\"}\"}'"
).read()
addr_list = [eval(addr) for addr in addr_str.split('\n') if addr]
self.logger.info("Scan address list:{0}".format(addr_list))
return addr_list
def _mc_instance(self, mac, state):
return MachineControl(
mac=mac,
state=MachineControlState(
ip=state.get('ip'),
profile_name=state.get('profile_name'),
is_capturing=state.get('is_capturing'),
is_shaping=state.get('is_shaping'),
online=state.get('online'),
capture_filter=state.get('capture_filter'),
last_update_time=state.get('last_update_time'),
))
def getMachineControls(self):
self.logger.info('Request getMachineControls')
addr_list = self._scanAddress()
new_machine_controls = {}
now = int(round(time.time() * 1000))
for addr in addr_list:
mac = addr["mac"]
ip = addr["ip"]
if self._machine_controls.get(mac):
# already has the mac addr
mc = self._machine_controls[mac]
mc['mac'] = mac
mc['ip'] = ip
mc['online'] = True
mc['last_update_time'] = now
# add to new_machine_controls & remove old mapping
new_machine_controls[mac] = mc
self._update_mcontrol(mc)
del self._machine_controls[mac]
else:
# add the shaping
mc = {
"mac": mac,
"ip": ip,
"is_capturing": False,
"is_shaping": False,
"online": True,
"last_update_time": now
}
new_machine_controls[mac] = mc
self._update_mcontrol(mc)
for mac in self._machine_controls:
self.unshapeMachine(mac)
mc = self._machine_controls[mac]
# check time out, limit 2 days
if now - mc['last_update_time'] > 2 * 24 * 60 * 60 * 1000:
self.db_task.queue.put(((mac), 'remove_mcontrol'))
else:
mc['online'] = False
mc['ip'] = ''
self._update_mcontrol(mc)
new_machine_controls[mac] = mc
self._machine_controls = new_machine_controls
self.logger.info("machine controls:{0}".format(self._machine_controls))
return [
self._mc_instance(mac, state)
for mac, state in self._machine_controls.items()
]
def updateMachineControl(self, update_mc):
if not self._machine_controls.get(update_mc.mac):
# can not update while not exist
return TrafficControlRc(
code=ReturnCode.INVALID_ADDRESS,
message="Invalid Address {mac:{0}, ip:{1}}".format(
update_mc.mac, update_mc.state.ip))
mc = self._machine_controls[update_mc.mac]
# update profile_name only for now
mc['profile_name'] = update_mc.state.profile_name
self._update_mcontrol(mc)
# update profiles while shaping
if mc['is_shaping']:
return self.shapeMachine(update_mc.mac)
return TrafficControlRc(code=ReturnCode.OK)
@address_check
def shapeMachine(self, mac, mc):
# remove old interface
if self._machine_shapings.get(mac):
# update shapings
old_shaping = self._machine_shapings[mac]
old_id = old_shaping['id']
self._unshape_interface(old_id, self.wan, old_shaping['ip'],
old_shaping['tc'].up)
self._unshape_interface(old_id, self.lan, old_shaping['ip'],
old_shaping['tc'].down)
self.idmanager.free(old_id)
del self._machine_shapings[mac]
# get profile setting
setting = self._profiles.get(mc.get('profile_name'))
if setting is None:
return TrafficControlRc(code=ACCESS_DENIED,
message="Invalid profile name: {0}".format(
mc['profile_name']))
new_id = None
try:
new_id = self.idmanager.new()
except Exception as e:
return TrafficControlRc(
code=ReturnCode.ID_EXHAUST,
message="No more session available: {0}".format(e))
self._machine_shapings[mac] = {
'id': new_id,
'ip': mc['ip'],
'tc': setting,
'profile_name': mc.get('profile_name')
}
# do shape
tcrc = self._shape_interface(new_id, self.wan, mc['ip'], setting.up)
if tcrc.code != ReturnCode.OK:
return tcrc
tcrc = self._shape_interface(new_id, self.lan, mc['ip'], setting.down)
if tcrc.code != ReturnCode.OK:
self._unshape_interface(new_id, self.wan, mc['ip'], setting.up)
return tcrc
mc['is_shaping'] = True
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
def unshapeMachine(self, mac):
# remove old interface
if self._machine_shapings.get(mac):
# update shapings
old_shaping = self._machine_shapings[mac]
old_id = old_shaping['id']
self._unshape_interface(old_id, self.wan, old_shaping['ip'],
old_shaping['tc'].up)
self._unshape_interface(old_id, self.lan, old_shaping['ip'],
old_shaping['tc'].down)
self.idmanager.free(old_id)
del self._machine_shapings[mac]
self._machine_controls[mac]['is_shaping'] = False
return TrafficControlRc(code=ReturnCode.OK)
def _unshape_interface(self, mark, eth, ip, settings):
"""Unshape traffic for a given IP/setting on a network interface
"""
raise NotImplementedError('Subclass should implement this')
def _shape_interface(self, mark, eth, ip, shaping):
"""Shape traffic for a given IP
"""
raise NotImplementedError('Subclass should implement this')
def _update_mcontrol(self, mc):
self.db_task.queue.put(
((mc['mac'], mc.get('ip'), mc.get('profile_name'),
mc['is_capturing'], mc['is_shaping'], mc['online'],
mc.get('capture_filter'), mc["last_update_time"]),
'add_mcontrol'))
@address_check
def getCapturePackets(self, mac, mc):
self.logger.info("getCapturePackets mac:{0}".format(mac))
pkts = self.scapy_task.getCapturePackets(mc['ip'])
if pkts is None:
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message="capture is not ready")
packet_dump = PacketsToJson(pkts)
# self.logger.info("packets: {}".format(packet_dump))
return TrafficControlRc(code=ReturnCode.OK, message=packet_dump)
@address_check
def startCapture(self, mac, mc, capture_filter):
self.logger.info("startCapture mac:{0} filter:{1}".format(
mac, capture_filter))
try:
self.scapy_task.startCapture(self.lan_name, capture_filter,
mc['ip'], mac)
except Exception as e:
self.logger.info("startCapture exception:{0}".format(str(e)))
mc['is_capturing'] = False
mc['capture_filter'] = None
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message=str(e))
mc['is_capturing'] = True
mc['capture_filter'] = capture_filter
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
@address_check
def stopCapture(self, mac, mc):
self.logger.info("stopCapture mac:{0}".format(mac))
self.scapy_task.stopCapture(self.lan_name, mc['ip'], mac)
mc['is_capturing'] = False
self._update_mcontrol(mc)
return TrafficControlRc(code=ReturnCode.OK)
@address_check
def exportPcap(self, mac, mc):
self.logger.info("exportPcap mac:{0}".format(mac))
pkts = self.scapy_task.getCapturePackets(mc['ip'])
if pkts is None:
return TrafficControlRc(code=ReturnCode.CAPTURE_NOT_READY,
message="capture is not ready")
if not os.path.exists(self.DEFAULT_PCAP_PATH):
os.makedirs(self.DEFAULT_PCAP_PATH)
pcap_path = self.DEFAULT_PCAP_PATH + "[" + mac + "].pcap"
wrpcap(pcap_path, pkts)
return TrafficControlRc(code=ReturnCode.OK)
class SetOptionTask(VTask):
LOOPLESS = True
some_option = option(type=int, default=0)
class NBServerTask(ThriftServerTask):
"""Spin up a thrift TNonblockingServer in a sparts worker thread"""
DEFAULT_HOST = '0.0.0.0'
DEFAULT_PORT = 0
OPT_PREFIX = 'thrift'
bound_host = bound_port = None
host = option(default=lambda cls: cls.DEFAULT_HOST, metavar='HOST',
help='Address to bind server to [%(default)s]')
port = option(default=lambda cls: cls.DEFAULT_PORT,
type=int, metavar='PORT',
help='Port to run server on [%(default)s]')
num_threads = option(name='threads', default=10, type=int, metavar='N',
help='Server Worker Threads [%(default)s]')
def initTask(self):
"""Overridden to bind sockets, etc"""
super(NBServerTask, self).initTask()
self._stopped = False
# Construct TServerSocket this way for compatibility with fbthrift
self.socket = TServerSocket(port=self.port)
self.socket.host = self.host
self.server = TNonblockingServer(self.processor, self.socket,
threads=self.num_threads)
self.server.prepare()
self.bound_addrs = []
for handle in self._get_socket_handles(self.server.socket):
addrinfo = handle.getsockname()
self.bound_host, self.bound_port = addrinfo[0:2]
self.logger.info("%s Server Started on %s", self.name,
self._fmt_hostport(self.bound_host, self.bound_port))
def _get_socket_handles(self, tsocket):
"""Helper to retrieve the socket objects for a given TServerSocket"""
handle = getattr(tsocket, 'handle', None)
if handle is not None:
return [tsocket.handle]
# Some TServerSocket implementations support multiple handles per
# TServerSocket (e.g., to support binding v4 and v6 without
# v4-mapped addresses
return tsocket.handles.values()
def _fmt_hostport(self, host, port):
if ':' in host:
return '[%s]:%d' % (host, port)
else:
return '%s:%d' % (host, port)
def stop(self):
"""Overridden to tell the thrift server to shutdown asynchronously"""
self.server.stop()
self.server.close()
self._stopped = True
def _runloop(self):
"""Overridden to execute TNonblockingServer's main loop"""
while not self.server._stop:
self.server.serve()
while not self._stopped:
time.sleep(0.1)
class PeriodicTask(VTask):
"""Task that executes `execute` at a specified interval
You must either override the `INTERVAL` (seconds) class attribute, or
pass a --{OPT_PREFIX}-interval in order for your task to run.
"""
INTERVAL = None
execute_duration_ms = samples(windows=[60, 240],
types=[SampleType.AVG, SampleType.MAX, SampleType.MIN])
n_iterations = counter()
n_slow_iterations = counter()
n_try_later = counter()
interval = option(type=float, metavar='SECONDS',
default=lambda cls: cls.INTERVAL,
help='How often this task should run [%(default)s] (s)')
def execute(self, context=None):
"""Override this to perform some custom action periodically."""
self.logger.debug('execute')
def execute_async(self):
f = Future()
if self.running:
# There's a race condition here. If the task has thrown but the
# thread(s) haven't stopped yet, you can enqueue a future that will
# never complete.
self.__futures.put(f)
else:
# If the task has stopped (e.g., due to a previous error),
# fail the future now and don't insert it into the queue.
f.set_exception(RuntimeError("Worker not running"))
return f
def has_pending(self):
return self.__futures.qsize() > 0
def initTask(self):
# Register an event that we can more smartly wait on in case shutdown
# is requested while we would be `sleep()`ing
self.stop_event = Event()
self.__futures = queue.Queue()
super(PeriodicTask, self).initTask()
assert self.interval is not None, \
"INTERVAL must be defined on %s or --%s-interval passed" % \
(self.name, self.name)
def stop(self):
self.stop_event.set()
super(PeriodicTask, self).stop()
def _runloop(self):
timer = Timer()
timer.start()
while not self.service._stop:
try:
result = self.execute()
# On a successful result, notify all blocked futures.
# Use pop like this to avoid race conditions.
while self.__futures.qsize():
f = self.__futures.get()
f.set_result(result)
except TryLater as e:
if self._handle_try_later(e):
return
continue
except Exception as e:
# On unhandled exceptions, set the exception on any async
# blocked execute calls.
while self.__futures.qsize():
f = self.__futures.get()
f.set_exception(e)
raise
self.n_iterations.increment()
self.execute_duration_ms.add(timer.elapsed * 1000)
to_sleep = self.interval - timer.elapsed
if to_sleep > 0:
if self.stop_event.wait(to_sleep):
return
else:
self.n_slow_iterations.increment()
timer.start()
def _handle_try_later(self, e):
self.n_try_later.increment()
if e.after is not None:
self.logger.debug("TryLater (%s) thrown. Retrying in %.2fs",
e.message, e.after)
else:
self.logger.debug("TryLater (%s) thrown. Retrying now",
e.message)
return self.stop_event.wait(e.after)
class QueueTask(VTask):
"""Task that calls `execute` for all work put on its `queue`"""
MAX_ITEMS = 0
WORKERS = 1
max_items = option(type=int,
default=lambda cls: cls.MAX_ITEMS,
help='Set a bounded queue length. This may '
'cause unexpected deadlocks. [%(default)s]')
workers = option(type=int,
default=lambda cls: cls.WORKERS,
help='Number of threads to spawn to work on items from '
'its queue. [%(default)s]')
def execute(self, item, context):
"""Implement this in your QueueTask subclasses"""
raise NotImplementedError()
def initTask(self):
super(QueueTask, self).initTask()
self.queue = queue.Queue(maxsize=self.max_items)
self._shutdown_sentinel = object()
def stop(self):
super(QueueTask, self).stop()
self.queue.put(self._shutdown_sentinel)
def _runloop(self):
while not self.service._stop:
try:
item = self.queue.get(timeout=1.0)
if item is self._shutdown_sentinel:
self.queue.put(item)
break
except queue.Empty:
continue
if isinstance(item, ExecuteContext):
context = item
item = context.item
else:
context = ExecuteContext(item=item)
try:
result = self.execute(item, context)
if context.deferred is not None:
context.deferred.callback(result)
except TryLater:
context.attempt += 1
self.queue.put(context)
except Exception as ex:
if context.deferred is not None:
self.unhandled = None
context.deferred.addErrback(self.unhandledErrback)
context.deferred.errback(ex)
if self.unhandled is not None:
self.unhandled.raiseException()
else:
raise
finally:
self.queue.task_done()
def unhandledErrback(self, error):
self.unhandled = error
return None
class MYSERVICE(VService):
basicopt = option(default="spam")
opt_uscore = option(default="eggs")
opt_uscore2 = option(name="opt_uscore2", default="ham")
class DBusServiceTask(DBusTask):
"""Glue Task for exporting this VService over DBus"""
OPT_PREFIX = 'dbus'
BUS_NAME = None
BUS_CLASS = VServiceDBusObject
USE_SYSTEM_BUS = False
bus_name = option(default=lambda cls: cls.BUS_NAME,
metavar='NAME',
help='Bus Name. Should be something like '
'"com.sparts.AwesomeService"')
replace = option(action='store_true',
type=bool,
default=False,
help='Replace, and enable replacing of this service')
queue = option(action='store_true',
type=bool,
default=False,
help='If not --{task}-replace, will wait to take '
'this bus name')
system_bus = option(action='store_true',
type=bool,
default=lambda cls: cls.USE_SYSTEM_BUS,
help='Use system bus')
dbus_service = None
def initTask(self):
super(DBusServiceTask, self).initTask()
assert self.bus_name is not None, \
"You must pass a --{task}-bus-name"
def _makeBus(self):
if self.system_bus:
return dbus.SystemBus(private=True)
return dbus.SessionBus(private=True)
def _asyncStartCb(self):
self.bus = self._makeBus()
self.dbus_service = dbus.service.BusName(self.bus_name, self.bus,
self.replace, self.replace,
self.queue)
return True
def _asyncStart(self):
res = self.asyncRun(self._asyncStartCb)
res.result()
def start(self):
self._asyncStart()
self.addHandlers()
super(DBusServiceTask, self).start()
def addHandlers(self):
self.sparts_dbus = self.BUS_CLASS(self)
if HAVE_FB303:
task = self.service.getTask(FB303HandlerTask)
if task is not None:
self.fb303_dbus = FB303DbusService(self.dbus_service, task,
self.service.name)
def _asyncStopCb(self):
self.dbus_service = None
# self.bus.close()
self.bus = None
return True
def _asyncStop(self):
res = self.asyncRun(self._asyncStopCb)
res.result()
def stop(self):
"""Run the bus cleanup code within the context of the main loop. The
task depends in DBusMainLoopTask, the stop() method will be called
before DBusMainLoopTask.stop() where loop.quit() is done.
"""
super(DBusServiceTask, self).stop()
self._asyncStop()
class TornadoHTTPTask(TornadoTask):
"""A loopless task that implements an HTTP server using Tornado.
It is loopless because it depends on tornado's separate IOLoop task. You
will need to subclass this to do something more useful."""
LOOPLESS = True
OPT_PREFIX = 'http'
DEFAULT_PORT = 0
DEFAULT_HOST = ''
DEFAULT_SOCK = ''
requests = counter()
#latency = samples(windows=[60, 3600],
# types=[SampleType.AVG, SampleType.MIN, SampleType.MAX])
host = option(metavar='HOST',
default=lambda cls: cls.DEFAULT_HOST,
help='Address to bind server to [%(default)s]')
port = option(metavar='PORT',
default=lambda cls: cls.DEFAULT_PORT,
help='Port to run server on [%(default)s]')
sock = option(metavar='PATH',
default=lambda cls: cls.DEFAULT_SOCK,
help='Default path to use for local file socket '
'[%(default)s]')
group = option(name='sock-group',
metavar='GROUP',
default='',
help='Group to create unix files as [%(default)s]')
def getApplicationConfig(self):
"""Override this to register custom handlers / routes."""
return [
('/', HelloWorldHandler),
]
def initTask(self):
super(TornadoHTTPTask, self).initTask()
self.app = tornado.web.Application(self.getApplicationConfig(),
log_function=self.tornadoRequestLog)
self.server = tornado.httpserver.HTTPServer(self.app)
if self.sock:
assert self.host == self.DEFAULT_HOST, \
"Do not specify host *and* sock (%s, %s)" % \
(self.host, self.sock)
assert int(self.port) == self.DEFAULT_PORT, \
"Do not specify port *and* sock (%s, %s)" % \
(self.port, self.DEFAULT_PORT)
gid, mode = -1, 0o600
if self.group != '':
e = grp.getgrnam(self.group)
gid, mode = e.gr_gid, 0o660
sock = tornado.netutil.bind_unix_socket(self.sock, mode=mode)
if gid != -1:
os.chown(self.sock, -1, gid)
self.server.add_sockets([sock])
else:
self.server.listen(self.port, self.host)
self.bound_addrs = []
for sock in itervalues(self.server._sockets):
sockaddr = sock.getsockname()
self.bound_addrs.append(sockaddr)
self.logger.info("%s Server Started on %s (port %s)", self.name,
sockaddr[0], sockaddr[1])
@property
def bound_v4_addrs(self):
return [a[0] for a in self.bound_addrs if len(a) == 2]
@property
def bound_v6_addrs(self):
return [a[0] for a in self.bound_addrs if len(a) == 4]
def tornadoRequestLog(self, handler):
self.requests.increment()
def stop(self):
super(TornadoHTTPTask, self).stop()
self.server.stop()
class TASK(VTask):
LOOPLESS = True
basicopt = option(default="spam")
opt_uscore = option(default="eggs")
opt_uscore2 = option(name="opt_uscore2", default="ham")
class DBusClientTask(DBusTask):
"""DBus client task helper class. Inherit this class and override
bus_service_online() and bus_service_offline() methods. Set
DBUS_SERVICE_NAME to the name of a DBus service that you want to
use. The class will automatically setup a name watch and call
online/offline callbacks when service becomes available.
"""
OPT_PREFIX = 'dbus-client'
DBUS_SERVICE_NAME = None
session_bus = option(action='store_true', type=bool,
default=False,
help='Use session bus to access service')
def __init__(self, *args, **kwargs):
super(DBusClientTask, self).__init__(*args, **kwargs)
self.bus = None
assert self.DBUS_SERVICE_NAME != None
def start(self):
self.logger.debug('starting client task for DBus service %s',
self.DBUS_SERVICE_NAME)
self.asyncRun(self._setup_bus_client)
super(DBusClientTask, self).start()
def _setup_bus_client(self):
"""Start task. Perform necessary setup:
- bus connection
- setup service name monitoring
"""
self.logger.debug('starting dbus client task for service: %s', self.service.name)
# get bus
if self.session_bus:
self.bus = dbus.SessionBus(private=True)
else:
self.bus = dbus.SystemBus(private=True)
self.logger.info('using %s bus for DBus service %s',
'session' if self.session_bus else 'system',
self.DBUS_SERVICE_NAME)
# we're monitoring service name changes
self.bus.watch_name_owner(self.DBUS_SERVICE_NAME,
self._bus_service_name_changed)
def stop(self):
"""Stop task"""
# get rid of reference to bus
self.bus = None
super(DBusClientTask, self).stop()
def _bus_service_name_changed(self, name):
"""Callback for when a name owner of bus service has changed
:param str name: service name, either actual service name or empty
"""
self.logger.debug('service name changed: %s', name)
if not name:
self.logger.info('service lost')
self.bus_service_offline()
else:
self.bus_service_online()
def bus_service_online(self):
"""Override this method to handle event when the service becomes avaialble"""
def bus_service_offline(self):
"""Override this method to handle event when the service becomes unavailable"""
def get_proxy(self, path, interface):
"""Obtain proxy to an interface `interface` defined in
DBUS_SERVICE_NAME service at path `path`
:param path str: object path
:param interface str: interface name
:return: interface proxy
"""
assert self.bus
try:
obj = self.bus.get_object(self.DBUS_SERVICE_NAME, path)
iface_proxy = dbus.Interface(obj, interface)
except dbus.DBusException:
self.logger.exception('failed to obtain proxy to servo')
return None
else:
return iface_proxy
class AtcdThriftHandlerTask(ThriftHandlerTask):
"""Atcd's thrift handler.
This is the main entry point of the program that implements the atcd.thrift
interface definition.
Platform specific behaviour will be implemented in Atcd`Platform`Shaper
class.
"""
ID_MANAGER_ID_MIN = 1
ID_MANAGER_ID_MAX = 2**16
MODULE = Atcd
DEPS = [AtcdDBQueueTask]
DEFAULT_LAN = 'eth1'
DEFAULT_WAN = 'eth0'
DEFAULT_IPTABLES = '/sbin/iptables'
DEFAULT_TCPDUMP = '/usr/sbin/tcpdump'
DEFAULT_PCAP_DIR = '/tmp'
DEFAULT_PCAP_URL_BASE = 'http://localhost:80'
DEFAULT_BURST_SIZE = 12000
DEFAULT_MODE = 'secure'
OPT_PREFIX = 'atcd'
lan_name = option(
default=DEFAULT_LAN,
metavar='LAN',
help='name of the LAN interface [%(default)s]',
name='lan',
)
wan_name = option(
default=DEFAULT_WAN,
metavar='WAN',
help='name of the WAN interface [%(default)s]',
name='wan',
)
iptables = option(default=DEFAULT_IPTABLES,
metavar='IPTABLES',
help='location of the iptables binary [%(default)s]')
tcpdump = option(default=DEFAULT_TCPDUMP,
metavar='TCPDUMP',
help='location of the tcpdump binary [%(default)s]')
pcap_dir = option(default=DEFAULT_PCAP_DIR,
metavar='PCAP_DIR',
help='Directory to store pcap files [%(default)s]')
pcap_url_base = option(default=DEFAULT_PCAP_URL_BASE,
metavar='PCAP_URL_BASE',
help='URL for pcap service [%(default)s]')
burst_size = option(
default=DEFAULT_BURST_SIZE,
metavar='BURST_SIZE',
type=int,
help='Amount of bytes that can be burst at a capped speed '
'[%(default)s]')
dont_drop_packets = option(
action='store_true',
help='[EXPERIMENTAL] Do not drop packets when going above max allowed'
' rate. Packets will be queued instead. Please mind that this'
' option will likely disappear in the future and is only provided'
' as a workaround until better longer term solution is found.',
)
fresh_start = option(
action='store_true',
help='Bypass saved shapings from a previous run [%(default)s]',
)
mode = option(
choices=['secure', 'unsecure'],
default=DEFAULT_MODE,
help='In which mode should atcd run? [%(default)s]',
)
@staticmethod
def factory():
"""Static method to discover and import the shaper to use.
Discover the platform on which Atcd is running and import the shaping
backend for this platform.
Returns:
The shaping backend class
Raises:
NotImplementedError: the shaping backend class couldn't be imported
"""
os_name = os.uname()[0]
klass = 'Atcd{0}Shaper'.format(os_name)
# If not imported yet, try to import
try:
if klass not in globals():
from_module_import_class(
'atcd.backends.{0}'.format(os_name.lower()), klass)
except AttributeError:
raise NotImplementedError('{0} is not implemented!'.format(klass))
except ImportError:
raise NotImplementedError('{0} backend is not implemented!'.format(
os_name.lower()))
return globals()[klass]
def initTask(self):
"""Thrift handler task initialization.
Performs the steps needed to initialize the shaping subsystem.
"""
super(AtcdThriftHandlerTask, self).initTask()
# Do this first because it can error out and it's better to
# error out before touching the networking stacks
self.db_task = self.service.tasks.AtcdDBQueueTask
self.lan = {'name': self.lan_name}
self.wan = {'name': self.wan_name}
self._links_lookup()
self._ip_to_id_map = {}
self._id_to_ip_map = {}
self.initialize_id_manager()
self.ip_to_pcap_proc_map = {}
self.initialize_shaping_system()
# Map of IP address to tc object that is currently
# being used to shape traffic from that IP.
# {ip: {'tc': tc, 'timeout': timeout}}
# {'10.0.0.1': {'tc': TrafficControl(...), 'timeout': 1234567890.1234}}
self._current_shapings = {}
self.access_manager = AccessManager(secure=self.mode != 'unsecure')
if not self.fresh_start:
self.logger.info('Restoring shaped connection from DB')
self._restore_saved_shapings()
def _links_lookup(self):
"""Initialize our mapping from network interface name to their device
id. Will raise and exception if one of the device is not found
"""
raise NotImplementedError('Subclass should implement this')
def initialize_id_manager(self):
"""Initialize the Id Manager. This is architecture dependant as the
shaping subsystems may have different requirements.
"""
self.idmanager = IdManager(first_id=type(self).ID_MANAGER_ID_MIN,
max_id=type(self).ID_MANAGER_ID_MAX)
def _restore_saved_shapings(self):
"""Restore the shapings from the sqlite3 db.
"""
# Isolate the things we are using eval on to reduce possible clownyness
# later on, also this way we don't have unused imports from importing
# blindly for eval
names = [
'TrafficControlledDevice', 'TrafficControl', 'Shaping',
'TrafficControlSetting', 'Loss', 'Delay', 'Corruption', 'Reorder'
]
globals = {name: getattr(atc_thrift.ttypes, name) for name in names}
# CurrentShapings(ip varchar primary key, tc blob, timeout int)
results = []
try:
results = self.db_task.get_saved_shapings()
except OperationalError:
self.logger.exception('Unable to perform DB operation')
for result in results:
tc = eval(result['tc'], globals)
timeout = float(result['timeout'])
if timeout > time.time():
tc.timeout = timeout - time.time()
try:
self.startShaping(tc)
except TrafficControlException as e:
# We have a shaping set in database that is denied
# probably because it was set in unsecure mode, passing
if (e.code == ReturnCode.ACCESS_DENIED
and self.mode == 'secure'):
self.logger.warn(
'Shaping Denied in secure mode, passing:'
' {0}'.format(e.message))
continue
raise
else:
self.db_task.queue.put(
(tc.device.controlledIP, 'remove_shaping'))
def stop(self):
"""Implements sparts.vtask.VTask.stop()
Each shaping platform should implement its own in order to clean
its state before shutting down the main loop.
"""
raise NotImplementedError('Subclass should implement this')
def initialize_shaping_system(self):
"""Initialize the shaping subsystem.
Each shaping platform should implement its own.
"""
raise NotImplementedError('Subclass should implement this')
def set_logger(self):
"""Initialize the logging subsystem.
"""
self.logger = logging.getLogger(__name__)
self.logger.setLevel(logging.DEBUG)
fmt = logging.Formatter(fmt=logging.BASIC_FORMAT)
# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(fmt=fmt)
self.logger.addHandler(ch)
# create syslog handler and set level to debug
sh = logging.handlers.SysLogHandler(address='/dev/log')
sh.setLevel(logging.DEBUG)
sh.setFormatter(fmt=fmt)
self.logger.addHandler(sh)
def getShapedDeviceCount(self):
"""Get the number of devices currently being shaped.
Returns:
The number of devices currently shaped.
"""
self.logger.info("Request getShapedDeviceCount")
return len(self._ip_to_id_map)
@AccessCheck
def startShaping(self, tc):
"""Start shaping a connection for a given device.
Implements the `startShaping` thrift method.
If the connection is already being shaped, the shaping will be updated
and the old one deleted.
Args:
A TrafficControl object that contains the device to be shaped, the
settings and the timeout.
Returns:
A TrafficControlRc object with code and message set to reflect
success/failure.
Raises:
A TrafficControlException with code and message set on uncaught
exception.
"""
self.logger.info("Request startShaping {0}".format(tc))
# Sanity checking
# IP
try:
socket.inet_aton(tc.device.controlledIP)
except Exception as e:
return TrafficControlRc(code=ReturnCode.INVALID_IP,
message="Invalid IP {}".format(
tc.device.controlledIP))
# timer
if tc.timeout < 0:
return TrafficControlRc(code=ReturnCode.INVALID_TIMEOUT,
message="Invalid Timeout {}".format(
tc.timeout))
new_id = None
try:
new_id = self.idmanager.new()
except Exception as e:
return TrafficControlRc(
code=ReturnCode.ID_EXHAUST,
message="No more session available: {0}".format(e))
old_id = self._ip_to_id_map.get(tc.device.controlledIP, None)
old_settings = self._current_shapings.get(tc.device.controlledIP,
{}).get('tc')
# start shaping for up way
for count, item in enumerate(tc.settings.up, start=1):
self.logger.info("up: start Shaping for item: {0}".format(item))
tcrc = self._shape_interface(
new_id + count,
self.wan,
tc.device.controlledIP,
item,
)
if tcrc.code != ReturnCode.OK:
return tcrc
# start shaping for down way
for count, item in enumerate(tc.settings.down, start=1):
self.logger.info("down: start Shaping for item: {0}".format(item))
tcrc = self._shape_interface(
new_id + count,
self.lan,
tc.device.controlledIP,
item,
)
# If we failed to set shaping for LAN interfaces, we should remove
# the shaping we just created for the WAN
if tcrc.code != ReturnCode.OK:
for count, item in enumerate(tc.settings.up, start=1):
self._unshape_interface(
old_id + count,
self.wan,
tc.device.controlledIP,
item,
)
return tcrc
self._add_mapping(new_id, tc)
self.db_task.queue.put(((tc, tc.timeout + time.time()), 'add_shaping'))
# if there were an existing id, remove it from dict
if old_id is not None:
# unshaping up way
for count, item in enumerate(old_settings.settings.up, start=1):
self._unshape_interface(
old_id + count,
self.wan,
tc.device.controlledIP,
item,
)
# unshaping down way
for count, item in enumerate(old_settings.settings.down, start=1):
self._unshape_interface(
old_id + count,
self.lan,
tc.device.controlledIP,
item,
)
del self._id_to_ip_map[old_id]
self.idmanager.free(old_id)
return TrafficControlRc(code=ReturnCode.OK)
@AccessCheck
def stopShaping(self, dev):
"""Stop shaping a connection for a given traffic controlled device.
Implements the `stopShaping` thrift method.
Args:
A TrafficControlledDevice object that contains the shaped device.
Returns:
A TrafficControlRc object with code and message set to reflect
success/failure.
Raises:
A TrafficControlException with code and message set on uncaught
exception.
"""
self.logger.info("Request stopShaping for ip {0}".format(
dev.controlledIP))
try:
socket.inet_aton(dev.controlledIP)
except Exception as e:
return TrafficControlRc(code=ReturnCode.INVALID_IP,
message="Invalid IP {0}: {1}".format(
dev.controlledIP, e))
id = self._ip_to_id_map.get(dev.controlledIP, None)
shaping = self._current_shapings.get(dev.controlledIP, {}).get('tc')
if id is not None:
# unshaping up way
for count, item in enumerate(shaping.settings.up, start=1):
self._unshape_interface(
id + count,
self.wan,
dev.controlledIP,
item,
)
# unshaping down way
for count, item in enumerate(shaping.settings.down, start=1):
self._unshape_interface(
id + count,
self.lan,
dev.controlledIP,
item,
)
self._del_mapping(id, dev.controlledIP)
self.db_task.queue.put((dev.controlledIP, 'remove_shaping'))
self.idmanager.free(id)
else:
return TrafficControlRc(
code=ReturnCode.UNKNOWN_SESSION,
message="No session for IP {} found".format(dev.controlledIP))
return TrafficControlRc(code=ReturnCode.OK)
def _unshape_interface(self, mark, eth, ip, settings):
"""Unshape traffic for a given IP/setting on a network interface
"""
raise NotImplementedError('Subclass should implement this')
def _shape_interface(self, mark, eth, ip, shaping):
"""Shape traffic for a given IP
"""
raise NotImplementedError('Subclass should implement this')
def isShaped(self, dev):
self.logger.info("Request isShaped for ip {0}".format(
dev.controlledIP))
return dev.controlledIP in self._ip_to_id_map
def getCurrentShaping(self, dev):
"""Get the TrafficControl object used to shape a
TrafficControlledDevice.
Args:
dev: a TrafficControlledDevice.
Returns:
A TrafficControl object representing the current shaping for the
device.
Raises:
A TrafficControlException if there is no TC object for that IP
"""
self.logger.info('Request getCurrentShaping for ip {0}'.format(
dev.controlledIP))
shaping = self._current_shapings.get(dev.controlledIP, {}).get('tc')
if shaping is None:
raise TrafficControlException(
code=ReturnCode.UNKNOWN_IP,
message='This IP ({0}) is not being shaped'.format(
dev.controlledIP))
return shaping
def _add_mapping(self, id, tc):
"""Adds a mapping from id to IP address and vice versa.
It also updates the dict mapping IPs to TrafficControl configs.
Args:
id: the id to map.
tc: the TrafficControl object to map.
"""
self._id_to_ip_map[id] = tc.device.controlledIP
self._ip_to_id_map[tc.device.controlledIP] = id
self._current_shapings[tc.device.controlledIP] = {
'tc': tc,
'timeout': time.time() + tc.timeout
}
def _del_mapping(self, id, ip):
"""Removes mappings from IP to id and id to IP.
Also remove the mapping from IP to TrafficControl configs.
"""
try:
del self._id_to_ip_map[id]
del self._ip_to_id_map[ip]
del self._current_shapings[ip]
except KeyError:
self.logger.exception('Unable to remove key from dict')
def run_cmd(self, cmd):
self.logger.info("Running {}".format(cmd))
return subprocess.call(shlex.split(cmd))
def _pcap_filename(self, ip, start_time):
return "%s_%d.cap" % (ip, start_time)
def _pcap_parse_filename(self, filename):
if filename.endswith(".cap"):
ip, start_time = filename.replace(".cap", "").split("_")
return ip, int(start_time)
def _pcap_url(self, filename):
return os.path.join(self.pcap_url_base, filename)
def _pcap_full_path(self, filename):
return os.path.join(self.pcap_dir, filename)
def _pcap_file_size(self, filename):
try:
return int(os.path.getsize(self._pcap_full_path(filename)))
except OSError:
return 0
def _cleanup_packet_capture_procs(self):
'''Delete finished procs from the map'''
for ip, p in self.ip_to_pcap_proc_map.items():
if not p or p.poll() is not None:
del self.ip_to_pcap_proc_map[ip]
@AccessCheck
def startPacketCapture(self, dev, timeout=3600):
"""Start a tcpdump process to capture packets for an ipaddr.
The process will run until the timeout expires or stopPacketCapture()
is called.
Args:
dev: a TrafficControlledDevice.
timeout: int Max time for tcpdump process to run.
Returns:
True if process started ok, otherwise False.
"""
self.logger.info(
"Request startPacketCapture for ip {0}, timeout {1}".format(
dev.controlledIP, timeout))
start_time = time.time()
filename = self._pcap_filename(dev.controlledIP, start_time)
cmd = """timeout {timeout!s}
{tcpdump} -vvv -s0 -i {eth} -w {filepath} host {ip}""".format(
timeout=timeout,
tcpdump=self.tcpdump,
eth=self.lan["name"],
filepath=self._pcap_full_path(filename),
ip=dev.controlledIP)
# Daemonize set the umask to 0o27 which prevents the http proxy service
# from reading the file. For lack of better solution for now, we can
# change the umask before spawning the subprocess and then restore its
# original value
umask = os.umask(0)
p = Popen(shlex.split(cmd))
os.umask(umask)
if p and p.poll() is None:
p.pcap = PacketCapture(ip=dev.controlledIP,
start_time=start_time,
file=PacketCaptureFile(
name=filename,
url=self._pcap_url(filename),
bytes=0),
pid=p.pid)
self.ip_to_pcap_proc_map[dev.controlledIP] = p
return p.pcap
else:
raise PacketCaptureException(
message="Failed to start tcpdump process")
@AccessCheck
def stopPacketCapture(self, dev):
"""Stop a tcpdump process that was started with startPacketCapture().
Args:
dev: a TrafficControlledDevice.
Returns:
The HTTP URL for the pcap file or empty string.
"""
self.logger.info("Request stopPacketCapture for ip {0}".format(
dev.controlledIP))
self._cleanup_packet_capture_procs()
if dev.controlledIP in self.ip_to_pcap_proc_map:
p = self.ip_to_pcap_proc_map[dev.controlledIP]
p.terminate()
# Wait a few secs for processes to die, while cleaning up dead ones
max_secs = 5
start_time = time.time()
while p.poll() is None and (time.time() - start_time) < max_secs:
time.sleep(0.5)
if p.poll() is None:
p.kill()
p.pcap.file.bytes = self._pcap_file_size(p.pcap.file.name)
return p.pcap
else:
raise PacketCaptureException(
message="No capture proc for given ipaddr")
def stopAllPacketCaptures(self):
"""Stop all running tcpdump procs.
"""
self.logger.info("Request stopAllPacketCaptures")
self._cleanup_packet_capture_procs()
if self.ip_to_pcap_proc_map:
for p in self.ip_to_pcap_proc_map.values():
p.terminate()
# Wait a few secs for processes to die, while cleaning up dead ones
max_secs = 5
start_time = time.time()
while self.ip_to_pcap_proc_map and \
(time.time() - start_time) < max_secs:
time.sleep(0.5)
self._cleanup_packet_capture_procs()
if self.ip_to_pcap_proc_map:
for p in self.ip_to_pcap_proc_map.values():
p.kill()
def listPacketCaptures(self, dev):
"""List the packet captures available for a given device.
Args:
dev: a TrafficControlledDevice.
Returns:
A list of PacketCapture ojbects.
"""
ip = dev.controlledIP
self.logger.info("Request listPacketCaptures for ip {0}".format(ip))
pcap_list = []
for filename in os.listdir(self.pcap_dir):
if not filename.endswith(".cap"):
continue
file_ip, start_time = self._pcap_parse_filename(filename)
if not file_ip == ip:
continue
pcap = PacketCapture(ip=ip,
start_time=start_time,
file=PacketCaptureFile(
name=filename,
url=self._pcap_url(filename),
bytes=self._pcap_file_size(filename)))
pcap_list.append(pcap)
return pcap_list
def listRunningPacketCaptures(self):
"""List the running packet captures.
Returns:
A list of PacketCapture ojbects.
"""
self.logger.info("Request listRunningPacketCaptures")
pcap_list = []
self._cleanup_packet_capture_procs()
for ip, p in self.ip_to_pcap_proc_map.items():
p.pcap.file.bytes = self._pcap_file_size(p.pcap.file.name)
pcap_list.append(p.pcap)
return pcap_list
def stop_expired_shapings(self):
"""Stop shaping that have expired.
"""
expired_devs = [
attrs['tc'].device
for ip, attrs in self._current_shapings.iteritems()
if attrs['timeout'] <= time.time()
]
for dev in expired_devs:
self.logger.info('Shaping for Device "{0}" expired'.format(dev))
self.logger.debug('calling stopShaping for "{0}"'.format(dev))
self.stopShaping(dev)
def requestToken(self, ip, duration):
"""Returns a unique, random access code.
Random token to be given to a host to control the `ip`.
The token validity is limited in time.
Args:
ip: The IP to control.
duration: How long the token will be valid for.
Returns:
An AccessToken.
"""
self.logger.info("Request requestToken({0}, {1})".format(ip, duration))
token = self.access_manager.generate_token(ip, duration)
return token
def requestRemoteControl(self, dev, accessToken):
"""Request to control a remote device.
Returns true if the token given is a valid token for the remote IP
according to the totp object stored for that IP
Args:
dev: The TrafficControlledDevice.
accessToken: The token to grant access.
Returns:
True if access is granted, False otherwise.
"""
self.logger.info("Request requestControl({0}, {1})".format(
dev, accessToken))
access_granted = False
try:
self.access_manager.validate_token(
dev,
accessToken,
)
access_granted = True
except AccessTokenException:
self.logger.exception("Access Denied for request")
return access_granted
def getDevicesControlledBy(self, ip):
"""Get the devices controlled by a given IP.
Args:
ip: The IP of the controlling host.
Returns:
A list of RemoteControlInstance.
"""
return self.access_manager.get_devices_controlled_by(ip)
def getDevicesControlling(self, ip):
"""Get the devices controlling a given IP.
Args:
ip: The IP of the controlled host.
Returns:
A list of RemoteControlInstance.
"""
return self.access_manager.get_devices_controlling(ip)
class CommandTask(TwistedTask):
"""A task that provides a useful API for executing other commands.
Python's Popen() can be hard to use, especially if you are executing long
running child processes, and need to handle various stdout, stderr, and
process exit events asynchronously.
This particular implementation relies on Twisted's ProcessProtocol, but it
wraps it in a way that makes it mostly opaque.
"""
LOOPLESS = True
OPT_PREFIX = 'cmd'
kill_timeout = option(type=float,
default=10.0,
help="Default shutdown kill timeout for outstanding "
"commands [%(default)s]")
started = counter()
finished = counter()
def run(self,
command,
on_stdout=None,
on_stderr=None,
on_exit=None,
line_buffered=True,
kill_timeout=None,
env=None):
"""Call this function to start a new child process running `command`.
Additional callbacks, such as `on_stdout`, `on_stderr`, and `on_exit`,
can be provided, that will receive a variety of parameters on the
appropriate events.
Line buffering can be disabled by passing `line_buffered`=False.
Also, a custom `kill_timeout` (seconds) may be set that overrides the
task default, in the event that a shutdown is received and you want
to allow more time for the command to shut down."""
self.logger.debug("task starting %s...", command)
if isinstance(command, six.string_types):
command = command.split(" ")
# wrap on_exit with helper to remove registered comments
on_exit = functools.partial(self._procExited, on_exit)
proto = _ProcessProtocolAdapter(on_stdout,
on_stderr,
on_exit,
line_buffered=line_buffered)
if twisted.python.threadable.isInIOThread():
result = self.reactor.spawnProcess(proto,
executable=command[0],
args=command)
else:
result = twisted.internet.threads.blockingCallFromThread(
self.reactor,
self.reactor.spawnProcess,
proto,
executable=command[0],
args=command,
env=env)
self.outstanding[result] = kill_timeout
self.started.increment()
return result
def initTask(self):
super(CommandTask, self).initTask()
self.outstanding = {}
def _procExited(self, on_exit, proto, trans, reason):
self.logger.debug("%s closed for %s", trans, reason)
if on_exit is not None:
on_exit(reason)
self.outstanding.pop(trans)
self.finished.increment()
return None
def join(self):
"""Overridden to block for process workers to shutdown / be killed."""
# TODO: Conditions instead of sleep polling?
while len(self.outstanding) > 0:
time.sleep(0.250)
def _killOutstanding(self, trans):
if trans in self.outstanding:
self.logger.info("Sending SIGKILL to %s", trans)
trans.signalProcess(signal.SIGKILL)
def stop(self):
# twisted is pretty smart; the default signal handlers it installs
# propagate SIGTERM to its children, so while we don't need to manually
# TERM, we might still need to set some kill timeouts
super(CommandTask, self).stop()
for trans, kill_timeout in self.outstanding.items():
if kill_timeout is None:
kill_timeout = self.kill_timeout
self.logger.info("Enqueuing kill for %s in %.1fs", trans,
kill_timeout)
args = (kill_timeout, self._killOutstanding, trans)
if twisted.python.threadable.isInIOThread():
self.reactor.callLater(*args)
else:
self.reactor.callFromThread(self.reactor.callLater, *args)
def isDoneWithReactor(self):
"""Overridden to keep reactor running until all commands finish."""
return len(self.outstanding) == 0
class AtcdDBQueueTestTask(AtcdDBQueueTask):
sqlite_file = option(default=atc_db_file().name)
class SetOptionTask(VTask):
LOOPLESS = True
some_option = option(type=int, default=0)
list_option = option(nargs='*', type=int, default=[])
other_list_option = option(nargs='*')
