def test_get_flow(monkeypatch):
proxy_processes = [
Mock(mode=ProxyMode.FORWARD),
Mock(mode=ProxyMode.REVERSE),
]
flow_state = {'mode': ProxyMode.FORWARD.value, 'timestamp_start': 1}
connections = [
(Mock(recv=Mock(return_value=flow_state)), Mock()),
(Mock(), Mock()),
]
monkeypatch.setattr(
'satellite.proxy.manager.ProxyProcess',
Mock(side_effect=proxy_processes),
)
monkeypatch.setattr(
'satellite.proxy.manager.Pipe',
Mock(side_effect=connections),
)
monkeypatch.setattr(
'satellite.proxy.manager.load_flow_from_state',
lambda state: Mock(**state),
)
flow_id = '23f11ab7-e071-4997-97f3-ace07bb9e56d'
manager = ProxyManager(9099, 9098, Mock())
manager._flows[flow_id] = ProxyMode.FORWARD
flow = manager.get_flow(flow_id)
assert flow.timestamp_start == 1
assert flow.mode == ProxyMode.FORWARD.value
connections[0][0].send.assert_called_once_with(
commands.GetFlowCommand(flow_id), )
def __init__(self, config: SatelliteConfig = None):
self.config = config or SatelliteConfig()
super().__init__(debug=self.config.debug)
self._should_exit = False
self.add_handlers(r'^(localhost|[0-9.]+|\[[0-9a-fA-F:]+\])$', [
(r'/', flow_handlers.Index),
(r'/updates', ClientConnection),
(r'/route', RoutesHandler),
(r'/route/(?P<route_id>[0-9a-f\-]+)', RouteHandler),
(r'/flows(?:\.json)?', flow_handlers.Flows),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)', flow_handlers.FlowHandler),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)/replay',
flow_handlers.ReplayFlow),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)/duplicate',
flow_handlers.DuplicateFlow),
(r'/logs/(?P<flow_id>[0-9a-f\-]+)',
audit_logs_handler.AuditLogsHandler),
])
self.proxy_manager = ProxyManager(
forward_proxy_port=self.config.forward_proxy_port,
reverse_proxy_port=self.config.reverse_proxy_port,
event_handler=partial(
self._proxy_event_handler,
loop=asyncio.get_event_loop(),
))
def test_get_flows(monkeypatch):
proxy_processes = [
Mock(mode=ProxyMode.FORWARD),
Mock(mode=ProxyMode.REVERSE),
]
connections = [
(Mock(recv=Mock(return_value=[{
'timestamp_start': 2
}])), Mock()),
(Mock(recv=Mock(return_value=[{
'timestamp_start': 1
}])), Mock()),
]
monkeypatch.setattr(
'satellite.proxy.manager.ProxyProcess',
Mock(side_effect=proxy_processes),
)
monkeypatch.setattr(
'satellite.proxy.manager.Pipe',
Mock(side_effect=connections),
)
monkeypatch.setattr(
'satellite.proxy.manager.load_flow_from_state',
lambda state: Mock(**state),
)
manager = ProxyManager(9099, 9098, Mock())
flows = manager.get_flows()
assert len(flows) == 2
assert flows[0].timestamp_start == 1
assert flows[1].timestamp_start == 2
for cmd_channel, _ in connections:
cmd_channel.send.assert_called_once_with(commands.GetFlowsCommand())
def test_update_flow(monkeypatch):
proxy_processes = [
Mock(mode=ProxyMode.FORWARD),
Mock(mode=ProxyMode.REVERSE),
]
connections = [
(Mock(), Mock()),
(Mock(), Mock()),
]
monkeypatch.setattr(
'satellite.proxy.manager.ProxyProcess',
Mock(side_effect=proxy_processes),
)
monkeypatch.setattr(
'satellite.proxy.manager.Pipe',
Mock(side_effect=connections),
)
flow_id = '23f11ab7-e071-4997-97f3-ace07bb9e56d'
flow_data = {'flow': 'data'}
manager = ProxyManager(9099, 9098, Mock())
manager._flows[flow_id] = ProxyMode.FORWARD
manager.update_flow(flow_id, flow_data)
connections[0][0].send.assert_called_once_with(
commands.UpdateFlowCommand(flow_id, flow_data), )
class WebApplication(Application):
def __init__(self, config: SatelliteConfig = None):
self.config = config or SatelliteConfig()
super().__init__(debug=self.config.debug)
self._should_exit = False
self.add_handlers(r'^(localhost|[0-9.]+|\[[0-9a-fA-F:]+\])$', [
(r'/', flow_handlers.Index),
(r'/updates', ClientConnection),
(r'/route', RoutesHandler),
(r'/route/(?P<route_id>[0-9a-f\-]+)', RouteHandler),
(r'/flows(?:\.json)?', flow_handlers.Flows),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)', flow_handlers.FlowHandler),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)/replay',
flow_handlers.ReplayFlow),
(r'/flows/(?P<flow_id>[0-9a-f\-]+)/duplicate',
flow_handlers.DuplicateFlow),
(r'/logs/(?P<flow_id>[0-9a-f\-]+)',
audit_logs_handler.AuditLogsHandler),
(r'/aliases', alias_handlers.AliasesHandler),
(r'/aliases/(?P<public_alias>[\w]+)', alias_handlers.AliasHandler),
])
self.proxy_manager = ProxyManager(
forward_proxy_port=self.config.forward_proxy_port,
reverse_proxy_port=self.config.reverse_proxy_port,
event_handler=partial(
self._proxy_event_handler,
loop=asyncio.get_event_loop(),
))
def _proxy_event_handler(self, event, loop):
asyncio.run_coroutine_threadsafe(
ClientConnection.process_proxy_event(event),
loop,
)
def start(self):
loop = asyncio.get_event_loop()
for sig in [signal.SIGINT, signal.SIGTERM]:
loop.add_signal_handler(sig, self.stop)
if self.settings.get('autoreload'):
autoreload.add_reload_hook(self.proxy_manager.stop)
self.proxy_manager.start()
self.listen(self.config.web_server_port)
logger.info(
f'Web server listening at {self.config.web_server_port} port.')
IOLoop.current().start()
def stop(self):
if self._should_exit:
return
self._should_exit = True
self.proxy_manager.stop()
IOLoop.current().stop()
def test_start_stop(monkeypatch):
proxy_processes = [
Mock(mode=ProxyMode.FORWARD),
Mock(mode=ProxyMode.REVERSE),
]
connections = [(Mock(), Mock()), (Mock(), Mock())]
monkeypatch.setattr(
'satellite.proxy.manager.ProxyProcess',
Mock(side_effect=proxy_processes),
)
monkeypatch.setattr(
'satellite.proxy.manager.Pipe',
Mock(side_effect=connections),
)
manager = ProxyManager(9099, 9098, Mock())
manager.start()
try:
for process in proxy_processes:
process.start.assert_called_once()
finally:
manager.stop()
for cmd_channel, _ in connections:
cmd_channel.send.assert_called_once_with(commands.StopCommand())
def test_audit_logs(monkeypatch):
proxy_processes = [
Mock(mode=ProxyMode.FORWARD),
Mock(mode=ProxyMode.REVERSE),
]
connections = [
(Mock(), Mock()),
(Mock(), Mock()),
]
make_proxy_process = Mock(side_effect=proxy_processes)
monkeypatch.setattr(
'satellite.proxy.manager.ProxyProcess',
make_proxy_process,
)
monkeypatch.setattr(
'satellite.proxy.manager.Pipe',
Mock(side_effect=connections),
)
manager = ProxyManager(9099, 9098, Mock())
manager.start()
try:
event_queue = make_proxy_process.call_args.kwargs['event_queue']
record = AuditLogTestRecord(
flow_id='flow-id',
proxy_mode=ProxyMode.FORWARD,
)
event_queue.put(
events.AuditLogEvent(
proxy_mode=ProxyMode.FORWARD,
record=record,
))
time.sleep(0.1)
assert manager.get_audit_logs('flow-id') == [record]
finally:
manager.stop()