diff options
| -rw-r--r-- | mitmproxy/controller.py | 40 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_controller.py | 7 |
3 files changed, 11 insertions, 38 deletions
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index c43fbb84..cc7e2e76 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -12,17 +12,18 @@ class ControlError(Exception): class Master(object): - """ - The master handles mitmproxy's main event loop. + The master handles mitmproxy's main event loop. """ - def __init__(self): self.event_queue = queue.Queue() self.should_exit = threading.Event() + self.servers = [] def start(self): self.should_exit.clear() + for server in self.servers: + ServerThread(server).start() def run(self): self.start() @@ -57,38 +58,18 @@ class Master(object): return changed def shutdown(self): + for server in self.servers: + server.shutdown() self.should_exit.set() - -class ServerMaster(Master): - - """ - The ServerMaster adds server thread support to the master. - """ - - def __init__(self): - super(ServerMaster, self).__init__() - self.servers = [] - def add_server(self, server): # We give a Channel to the server which can be used to communicate with the master channel = Channel(self.event_queue, self.should_exit) server.set_channel(channel) self.servers.append(server) - def start(self): - super(ServerMaster, self).start() - for server in self.servers: - ServerThread(server).start() - - def shutdown(self): - for server in self.servers: - server.shutdown() - super(ServerMaster, self).shutdown() - class ServerThread(threading.Thread): - def __init__(self, server): self.server = server super(ServerThread, self).__init__() @@ -100,12 +81,10 @@ class ServerThread(threading.Thread): class Channel(object): - """ - The only way for the proxy server to communicate with the master - is to use the channel it has been given. + The only way for the proxy server to communicate with the master + is to use the channel it has been given. """ - def __init__(self, q, should_exit): self.q = q self.should_exit = should_exit @@ -141,12 +120,10 @@ class Channel(object): class DummyReply(object): - """ A reply object that does nothing. Useful when we need an object to seem like it has a channel, and during testing. """ - def __init__(self): self.acked = False self.taken = False @@ -192,7 +169,6 @@ def handler(f): class Reply(object): - """ Messages sent through a channel are decorated with a "reply" attribute. This object is used to respond to the message through the return diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 1b4a999a..716d9a17 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -638,7 +638,7 @@ class State(object): self.flows.kill_all(master) -class FlowMaster(controller.ServerMaster): +class FlowMaster(controller.Master): @property def server(self): diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index c9a8e2f4..fbbebaa6 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -16,7 +16,6 @@ class TMsg: class TestMaster(object): def test_simple(self): - class DummyMaster(controller.Master): @controller.handler def handle_panic(self, _): @@ -34,10 +33,8 @@ class TestMaster(object): m.run() assert m.should_exit.is_set() - -class TestServerMaster(object): - def test_simple(self): - m = controller.ServerMaster() + def test_server_simple(self): + m = controller.Master() s = DummyServer(None) m.add_server(s) m.start() |