diff options
| -rw-r--r-- | mitmproxy/flow.py | 61 | ||||
| -rw-r--r-- | mitmproxy/models/__init__.py | 5 | ||||
| -rw-r--r-- | mitmproxy/models/flow.py | 9 | ||||
| -rw-r--r-- | mitmproxy/models/http.py | 6 | ||||
| -rw-r--r-- | mitmproxy/models/tcp.py | 50 | ||||
| -rw-r--r-- | mitmproxy/protocol/rawtcp.py | 61 | ||||
| -rw-r--r-- | mitmproxy/proxy/root_context.py | 2 |
7 files changed, 139 insertions, 55 deletions
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index ccedd1d4..1d05d4bb 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -3,7 +3,6 @@ """ from __future__ import absolute_import -import traceback from abc import abstractmethod, ABCMeta import hashlib import sys @@ -18,12 +17,13 @@ from typing import List, Optional, Set from netlib import wsgi, odict from netlib.exceptions import HttpException from netlib.http import Headers, http1, cookies +from netlib.utils import clean_bin from . import controller, tnetstring, filt, script, version, flow_format_compat from .onboarding import app from .proxy.config import HostMatcher from .protocol.http_replay import RequestReplayThread from .exceptions import Kill, FlowReadException -from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES +from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES, TCPFlow from collections import defaultdict @@ -900,6 +900,17 @@ class FlowMaster(controller.ServerMaster): self.handle_response(f) if f.error: self.handle_error(f) + elif isinstance(f, TCPFlow): + messages = f.messages + f.messages = [] + f.reply = controller.DummyReply() + self.handle_tcp_open(f) + while messages: + f.messages.append(messages.pop(0)) + self.handle_tcp_message(f) + if f.error: + self.handle_tcp_error(f) + self.handle_tcp_close(f) else: raise NotImplementedError() @@ -1087,18 +1098,52 @@ class FlowMaster(controller.ServerMaster): self.add_event('"{}" reloaded.'.format(s.filename), 'info') return ok - def handle_tcp_message(self, m): - self.run_script_hook("tcp_message", m) - m.reply() + def handle_tcp_open(self, flow): + self.state.add_flow(flow) + self.run_script_hook("tcp_open", flow) + flow.reply() + + def handle_tcp_message(self, flow): + self.run_script_hook("tcp_message", flow) + message = flow.messages[-1] + direction = "->" if message.from_client else "<-" + self.add_event("{client} {direction} tcp {direction} {server}".format( + client=repr(flow.client_conn.address), + server=repr(flow.server_conn.address), + direction=direction, + ), "info") + self.add_event(clean_bin(message.content), "debug") + flow.reply() + + def handle_tcp_error(self, flow): + if self.stream: + self.stream.add(flow) + self.add_event("Error in TCP connection to {}: {}".format( + repr(flow.server_conn.address), + flow.error + ), "info") + self.run_script_hook("tcp_error", flow) + flow.reply() + + def handle_tcp_close(self, flow): + self.state.delete_flow(flow) + if self.stream: + self.stream.add(flow) + self.run_script_hook("tcp_close", flow) + flow.reply() def shutdown(self): super(FlowMaster, self).shutdown() # Add all flows that are still active if self.stream: - for i in self.state.flows: - if not i.response: - self.stream.add(i) + for flow in self.state.flows: + # FIXME: We actually need to keep track of which flows are still active. + if isinstance(flow, HTTPFlow) and not flow.response: + self.stream.add(flow) + if isinstance(flow, TCPFlow): + # (assuming mitmdump only, this must be still active) + self.stream.add(flow) self.stop_stream() self.unload_scripts() diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py index df86eff4..3d9d9dae 100644 --- a/mitmproxy/models/__init__.py +++ b/mitmproxy/models/__init__.py @@ -7,9 +7,11 @@ from .http import ( from netlib.http import decoded from .connections import ClientConnection, ServerConnection from .flow import Flow, Error +from .tcp import TCPFlow FLOW_TYPES = dict( - http=HTTPFlow + http=HTTPFlow, + tcp=TCPFlow, ) __all__ = [ @@ -18,5 +20,6 @@ __all__ = [ "make_connect_response", "expect_continue_response", "ClientConnection", "ServerConnection", "Flow", "Error", + "TCPFlow" "FLOW_TYPES" ] diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py index 594147ec..1019c9fb 100644 --- a/mitmproxy/models/flow.py +++ b/mitmproxy/models/flow.py @@ -40,6 +40,9 @@ class Error(stateobject.StateObject): def __str__(self): return self.msg + def __repr__(self): + return self.msg + @classmethod def from_state(cls, state): # the default implementation assumes an empty constructor. Override @@ -99,6 +102,12 @@ class Flow(stateobject.StateObject): self._backup = state.pop("backup") super(Flow, self).set_state(state) + @classmethod + def from_state(cls, state): + f = cls(None, None) + f.set_state(state) + return f + def copy(self): f = copy.copy(self) diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py index 77a809cf..75ffbfd0 100644 --- a/mitmproxy/models/http.py +++ b/mitmproxy/models/http.py @@ -191,12 +191,6 @@ class HTTPFlow(Flow): response=HTTPResponse ) - @classmethod - def from_state(cls, state): - f = cls(None, None) - f.set_state(state) - return f - def __repr__(self): s = "<HTTPFlow" for a in ("request", "response", "error", "client_conn", "server_conn"): diff --git a/mitmproxy/models/tcp.py b/mitmproxy/models/tcp.py new file mode 100644 index 00000000..22fd1f0c --- /dev/null +++ b/mitmproxy/models/tcp.py @@ -0,0 +1,50 @@ +import time +from typing import List + +from netlib.utils import Serializable +from .flow import Flow + + +class TCPMessage(Serializable): + def __init__(self, from_client, content, timestamp=None): + self.content = content + self.from_client = from_client + if timestamp is None: + timestamp = time.time() + self.timestamp = timestamp + + @classmethod + def from_state(cls, state): + return cls(*state) + + def get_state(self): + return self.from_client, self.content, self.timestamp + + def set_state(self, state): + self.from_client = state.pop("from_client") + self.content = state.pop("content") + self.timestamp = state.pop("timestamp") + + def __repr__(self): + return "{direction} {content}".format( + direction="->" if self.from_client else "<-", + content=repr(self.content) + ) + + +class TCPFlow(Flow): + """ + A SSHFlow is a simplified representation of an SSH session. + """ + + def __init__(self, client_conn, server_conn, live=None): + super(TCPFlow, self).__init__("tcp", client_conn, server_conn, live) + self.messages = [] # type: List[TCPMessage] + + _stateobject_attributes = Flow._stateobject_attributes.copy() + _stateobject_attributes.update( + messages=List[TCPMessage] + ) + + def __repr__(self): + return "<TCPFlow ({} messages)>".format(len(self.messages)) diff --git a/mitmproxy/protocol/rawtcp.py b/mitmproxy/protocol/rawtcp.py index 7d18025e..5f6fca75 100644 --- a/mitmproxy/protocol/rawtcp.py +++ b/mitmproxy/protocol/rawtcp.py @@ -9,29 +9,26 @@ from netlib.exceptions import TcpException from netlib.tcp import ssl_read_select from netlib.utils import clean_bin from ..exceptions import ProtocolException -from .base import Layer - +from ..models import Error +from ..models.tcp import TCPFlow, TCPMessage -class TcpMessage(object): - - def __init__(self, client_conn, server_conn, sender, receiver, message): - self.client_conn = client_conn - self.server_conn = server_conn - self.sender = sender - self.receiver = receiver - self.message = message +from .base import Layer class RawTCPLayer(Layer): chunk_size = 4096 - def __init__(self, ctx, logging=True): - self.logging = logging + def __init__(self, ctx, ignore=False): + self.ignore = ignore super(RawTCPLayer, self).__init__(ctx) def __call__(self): self.connect() + if not self.ignore: + flow = TCPFlow(self.client_conn, self.server_conn, self) + self.channel.ask("tcp_open", flow) + buf = memoryview(bytearray(self.chunk_size)) client = self.client_conn.connection @@ -51,38 +48,24 @@ class RawTCPLayer(Layer): if isinstance(conn, SSL.Connection): # We can't half-close a connection, so we just close everything here. # Sockets will be cleaned up on a higher level. - return + break else: dst.shutdown(socket.SHUT_WR) if len(conns) == 0: - return + break continue - tcp_message = TcpMessage( - self.client_conn, self.server_conn, - self.client_conn if dst == server else self.server_conn, - self.server_conn if dst == server else self.client_conn, - buf[:size].tobytes()) - self.channel.ask("tcp_message", tcp_message) - dst.sendall(tcp_message.message) - - if self.logging: - # log messages are prepended with the client address, - # hence the "weird" direction string. - if dst == server: - direction = "-> tcp -> {}".format(repr(self.server_conn.address)) - else: - direction = "<- tcp <- {}".format(repr(self.server_conn.address)) - data = clean_bin(tcp_message.message) - self.log( - "{}\r\n{}".format(direction, data), - "info" - ) + tcp_message = TCPMessage(dst == server, buf[:size].tobytes()) + if not self.ignore: + flow.messages.append(tcp_message) + self.channel.ask("tcp_message", flow) + dst.sendall(tcp_message.content) except (socket.error, TcpException, SSL.Error) as e: - six.reraise( - ProtocolException, - ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))), - sys.exc_info()[2] - ) + if not self.ignore: + flow.error = Error("TCP connection closed unexpectedly: {}".format(repr(e))) + self.channel.tell("tcp_error", flow) + finally: + if not self.ignore: + self.channel.tell("tcp_close", flow) diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py index 9caae02a..21478dfd 100644 --- a/mitmproxy/proxy/root_context.py +++ b/mitmproxy/proxy/root_context.py @@ -65,7 +65,7 @@ class RootContext(object): else: ignore = self.config.check_ignore((client_hello.client_sni, 443)) if ignore: - return RawTCPLayer(top_layer, logging=False) + return RawTCPLayer(top_layer, ignore=True) # 2. Always insert a TLS layer, even if there's neither client nor server tls. # An inline script may upgrade from http to https, |