diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | mitmproxy/io/__init__.py | 4 | ||||
| -rw-r--r-- | mitmproxy/io/db.py | 9 | ||||
| -rw-r--r-- | mitmproxy/io/proto/http.proto | 1 | ||||
| -rw-r--r-- | mitmproxy/io/proto/http_pb2.py | 43 | ||||
| -rw-r--r-- | mitmproxy/io/protobuf.py | 53 | ||||
| -rw-r--r-- | setup.cfg | 2 | ||||
| -rw-r--r-- | setup.py | 2 | ||||
| -rwxr-xr-x | test/filename_matching.py | 3 | ||||
| -rw-r--r-- | test/mitmproxy/io/test_db.py | 26 | ||||
| -rw-r--r-- | test/mitmproxy/io/test_protobuf.py | 120 |
11 files changed, 200 insertions, 64 deletions
@@ -5,6 +5,7 @@ MANIFEST *.py[cdo] *.swp *.swo +*.sqlite *.egg-info/ .coverage* .idea diff --git a/mitmproxy/io/__init__.py b/mitmproxy/io/__init__.py index 854d2505..bd248fd2 100644 --- a/mitmproxy/io/__init__.py +++ b/mitmproxy/io/__init__.py @@ -1,8 +1,8 @@ from .io import FlowWriter, FlowReader, FilteredFlowWriter, read_flows_from_paths -from .db import DbHandler +from .db import DBHandler __all__ = [ - "FlowWriter", "FlowReader", "FilteredFlowWriter", "read_flows_from_paths", "DbHandler" + "FlowWriter", "FlowReader", "FilteredFlowWriter", "read_flows_from_paths", "DBHandler" ] diff --git a/mitmproxy/io/db.py b/mitmproxy/io/db.py index 01d30aa2..ea42e08b 100644 --- a/mitmproxy/io/db.py +++ b/mitmproxy/io/db.py @@ -2,17 +2,18 @@ import sqlite3 import os from mitmproxy.io import protobuf -from mitmproxy.http import HTTPFlow -from mitmproxy import exceptions -class DbHandler: +class DBHandler: """ This class is wrapping up connection to SQLITE DB. """ - def __init__(self, db_path="/tmp/tmp.sqlite"): + def __init__(self, db_path, mode='load'): + if mode == 'write': + if os.path.isfile(db_path): + os.remove(db_path) self.db_path = db_path self._con = sqlite3.connect(self.db_path) self._c = self._con.cursor() diff --git a/mitmproxy/io/proto/http.proto b/mitmproxy/io/proto/http.proto index 6c372433..c86a04f3 100644 --- a/mitmproxy/io/proto/http.proto +++ b/mitmproxy/io/proto/http.proto @@ -10,7 +10,6 @@ message HTTPFlow { optional bool marked = 7; optional string mode = 8; optional string id = 9; - optional int32 version = 10; } message HTTPRequest { diff --git a/mitmproxy/io/proto/http_pb2.py b/mitmproxy/io/proto/http_pb2.py index 9809e7ee..47b09815 100644 --- a/mitmproxy/io/proto/http_pb2.py +++ b/mitmproxy/io/proto/http_pb2.py @@ -19,7 +19,7 @@ DESCRIPTOR = _descriptor.FileDescriptor( name='http.proto', package='', syntax='proto2', - serialized_pb=_b('\n\nhttp.proto\"\x85\x02\n\x08HTTPFlow\x12\x1d\n\x07request\x18\x01 \x01(\x0b\x32\x0c.HTTPRequest\x12\x1f\n\x08response\x18\x02 \x01(\x0b\x32\r.HTTPResponse\x12\x19\n\x05\x65rror\x18\x03 \x01(\x0b\x32\n.HTTPError\x12&\n\x0b\x63lient_conn\x18\x04 \x01(\x0b\x32\x11.ClientConnection\x12&\n\x0bserver_conn\x18\x05 \x01(\x0b\x32\x11.ServerConnection\x12\x13\n\x0bintercepted\x18\x06 \x01(\x08\x12\x0e\n\x06marked\x18\x07 \x01(\x08\x12\x0c\n\x04mode\x18\x08 \x01(\t\x12\n\n\x02id\x18\t \x01(\t\x12\x0f\n\x07version\x18\n \x01(\x05\"\xfa\x01\n\x0bHTTPRequest\x12\x19\n\x11\x66irst_line_format\x18\x01 \x01(\t\x12\x0e\n\x06method\x18\x02 \x01(\t\x12\x0e\n\x06scheme\x18\x03 \x01(\t\x12\x0c\n\x04host\x18\x04 \x01(\t\x12\x0c\n\x04port\x18\x05 \x01(\x05\x12\x0c\n\x04path\x18\x06 \x01(\t\x12\x14\n\x0chttp_version\x18\x07 \x01(\t\x12\x1c\n\x07headers\x18\x08 \x03(\x0b\x32\x0b.HTTPHeader\x12\x0f\n\x07\x63ontent\x18\t \x01(\x0c\x12\x17\n\x0ftimestamp_start\x18\n \x01(\x01\x12\x15\n\rtimestamp_end\x18\x0b \x01(\x01\x12\x11\n\tis_replay\x18\x0c \x01(\x08\"\xbb\x01\n\x0cHTTPResponse\x12\x14\n\x0chttp_version\x18\x01 \x01(\t\x12\x13\n\x0bstatus_code\x18\x02 \x01(\x05\x12\x0e\n\x06reason\x18\x03 \x01(\t\x12\x1c\n\x07headers\x18\x04 \x03(\x0b\x32\x0b.HTTPHeader\x12\x0f\n\x07\x63ontent\x18\x05 \x01(\x0c\x12\x17\n\x0ftimestamp_start\x18\x06 \x01(\x01\x12\x15\n\rtimestamp_end\x18\x07 \x01(\x01\x12\x11\n\tis_replay\x18\x08 \x01(\x08\"+\n\tHTTPError\x12\x0b\n\x03msg\x18\x01 \x01(\t\x12\x11\n\ttimestamp\x18\x02 \x01(\x01\")\n\nHTTPHeader\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t\"%\n\x07\x41\x64\x64ress\x12\x0c\n\x04host\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\"\xc2\x02\n\x10\x43lientConnection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x19\n\x07\x61\x64\x64ress\x18\x02 \x01(\x0b\x32\x08.Address\x12\x17\n\x0ftls_established\x18\x03 \x01(\x08\x12\x12\n\nclientcert\x18\x04 \x01(\t\x12\x10\n\x08mitmcert\x18\x05 \x01(\t\x12\x17\n\x0ftimestamp_start\x18\x06 \x01(\x01\x12\x1b\n\x13timestamp_tls_setup\x18\x07 \x01(\x01\x12\x15\n\rtimestamp_end\x18\x08 \x01(\x01\x12\x0b\n\x03sni\x18\t \x01(\t\x12\x13\n\x0b\x63ipher_name\x18\n \x01(\t\x12\x1d\n\x15\x61lpn_proto_negotiated\x18\x0b \x01(\x0c\x12\x13\n\x0btls_version\x18\x0c \x01(\t\x12%\n\x0etls_extensions\x18\r \x03(\x0b\x32\r.TLSExtension\"\xeb\x02\n\x10ServerConnection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x19\n\x07\x61\x64\x64ress\x18\x02 \x01(\x0b\x32\x08.Address\x12\x1c\n\nip_address\x18\x03 \x01(\x0b\x32\x08.Address\x12 \n\x0esource_address\x18\x04 \x01(\x0b\x32\x08.Address\x12\x17\n\x0ftls_established\x18\x05 \x01(\x08\x12\x0c\n\x04\x63\x65rt\x18\x06 \x01(\t\x12\x0b\n\x03sni\x18\x07 \x01(\t\x12\x1d\n\x15\x61lpn_proto_negotiated\x18\x08 \x01(\x0c\x12\x13\n\x0btls_version\x18\t \x01(\t\x12\x17\n\x0ftimestamp_start\x18\n \x01(\x01\x12\x1b\n\x13timestamp_tcp_setup\x18\x0b \x01(\x01\x12\x1b\n\x13timestamp_tls_setup\x18\x0c \x01(\x01\x12\x15\n\rtimestamp_end\x18\r \x01(\x01\x12\x1e\n\x03via\x18\x0e \x01(\x0b\x32\x11.ServerConnection\"*\n\x0cTLSExtension\x12\x0b\n\x03int\x18\x01 \x01(\x03\x12\r\n\x05\x62ytes\x18\x02 \x01(\x0c') + serialized_pb=_b('\n\nhttp.proto\"\xf4\x01\n\x08HTTPFlow\x12\x1d\n\x07request\x18\x01 \x01(\x0b\x32\x0c.HTTPRequest\x12\x1f\n\x08response\x18\x02 \x01(\x0b\x32\r.HTTPResponse\x12\x19\n\x05\x65rror\x18\x03 \x01(\x0b\x32\n.HTTPError\x12&\n\x0b\x63lient_conn\x18\x04 \x01(\x0b\x32\x11.ClientConnection\x12&\n\x0bserver_conn\x18\x05 \x01(\x0b\x32\x11.ServerConnection\x12\x13\n\x0bintercepted\x18\x06 \x01(\x08\x12\x0e\n\x06marked\x18\x07 \x01(\x08\x12\x0c\n\x04mode\x18\x08 \x01(\t\x12\n\n\x02id\x18\t \x01(\t\"\xfa\x01\n\x0bHTTPRequest\x12\x19\n\x11\x66irst_line_format\x18\x01 \x01(\t\x12\x0e\n\x06method\x18\x02 \x01(\t\x12\x0e\n\x06scheme\x18\x03 \x01(\t\x12\x0c\n\x04host\x18\x04 \x01(\t\x12\x0c\n\x04port\x18\x05 \x01(\x05\x12\x0c\n\x04path\x18\x06 \x01(\t\x12\x14\n\x0chttp_version\x18\x07 \x01(\t\x12\x1c\n\x07headers\x18\x08 \x03(\x0b\x32\x0b.HTTPHeader\x12\x0f\n\x07\x63ontent\x18\t \x01(\x0c\x12\x17\n\x0ftimestamp_start\x18\n \x01(\x01\x12\x15\n\rtimestamp_end\x18\x0b \x01(\x01\x12\x11\n\tis_replay\x18\x0c \x01(\x08\"\xbb\x01\n\x0cHTTPResponse\x12\x14\n\x0chttp_version\x18\x01 \x01(\t\x12\x13\n\x0bstatus_code\x18\x02 \x01(\x05\x12\x0e\n\x06reason\x18\x03 \x01(\t\x12\x1c\n\x07headers\x18\x04 \x03(\x0b\x32\x0b.HTTPHeader\x12\x0f\n\x07\x63ontent\x18\x05 \x01(\x0c\x12\x17\n\x0ftimestamp_start\x18\x06 \x01(\x01\x12\x15\n\rtimestamp_end\x18\x07 \x01(\x01\x12\x11\n\tis_replay\x18\x08 \x01(\x08\"+\n\tHTTPError\x12\x0b\n\x03msg\x18\x01 \x01(\t\x12\x11\n\ttimestamp\x18\x02 \x01(\x01\")\n\nHTTPHeader\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t\"%\n\x07\x41\x64\x64ress\x12\x0c\n\x04host\x18\x01 \x01(\t\x12\x0c\n\x04port\x18\x02 \x01(\x05\"\xc2\x02\n\x10\x43lientConnection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x19\n\x07\x61\x64\x64ress\x18\x02 \x01(\x0b\x32\x08.Address\x12\x17\n\x0ftls_established\x18\x03 \x01(\x08\x12\x12\n\nclientcert\x18\x04 \x01(\t\x12\x10\n\x08mitmcert\x18\x05 \x01(\t\x12\x17\n\x0ftimestamp_start\x18\x06 \x01(\x01\x12\x1b\n\x13timestamp_tls_setup\x18\x07 \x01(\x01\x12\x15\n\rtimestamp_end\x18\x08 \x01(\x01\x12\x0b\n\x03sni\x18\t \x01(\t\x12\x13\n\x0b\x63ipher_name\x18\n \x01(\t\x12\x1d\n\x15\x61lpn_proto_negotiated\x18\x0b \x01(\x0c\x12\x13\n\x0btls_version\x18\x0c \x01(\t\x12%\n\x0etls_extensions\x18\r \x03(\x0b\x32\r.TLSExtension\"\xeb\x02\n\x10ServerConnection\x12\n\n\x02id\x18\x01 \x01(\t\x12\x19\n\x07\x61\x64\x64ress\x18\x02 \x01(\x0b\x32\x08.Address\x12\x1c\n\nip_address\x18\x03 \x01(\x0b\x32\x08.Address\x12 \n\x0esource_address\x18\x04 \x01(\x0b\x32\x08.Address\x12\x17\n\x0ftls_established\x18\x05 \x01(\x08\x12\x0c\n\x04\x63\x65rt\x18\x06 \x01(\t\x12\x0b\n\x03sni\x18\x07 \x01(\t\x12\x1d\n\x15\x61lpn_proto_negotiated\x18\x08 \x01(\x0c\x12\x13\n\x0btls_version\x18\t \x01(\t\x12\x17\n\x0ftimestamp_start\x18\n \x01(\x01\x12\x1b\n\x13timestamp_tcp_setup\x18\x0b \x01(\x01\x12\x1b\n\x13timestamp_tls_setup\x18\x0c \x01(\x01\x12\x15\n\rtimestamp_end\x18\r \x01(\x01\x12\x1e\n\x03via\x18\x0e \x01(\x0b\x32\x11.ServerConnection\"*\n\x0cTLSExtension\x12\x0b\n\x03int\x18\x01 \x01(\x03\x12\r\n\x05\x62ytes\x18\x02 \x01(\x0c') ) @@ -95,13 +95,6 @@ _HTTPFLOW = _descriptor.Descriptor( message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='version', full_name='HTTPFlow.version', index=9, - number=10, type=5, cpp_type=1, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - options=None, file=DESCRIPTOR), ], extensions=[ ], @@ -115,7 +108,7 @@ _HTTPFLOW = _descriptor.Descriptor( oneofs=[ ], serialized_start=15, - serialized_end=276, + serialized_end=259, ) @@ -222,8 +215,8 @@ _HTTPREQUEST = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=279, - serialized_end=529, + serialized_start=262, + serialized_end=512, ) @@ -302,8 +295,8 @@ _HTTPRESPONSE = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=532, - serialized_end=719, + serialized_start=515, + serialized_end=702, ) @@ -340,8 +333,8 @@ _HTTPERROR = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=721, - serialized_end=764, + serialized_start=704, + serialized_end=747, ) @@ -378,8 +371,8 @@ _HTTPHEADER = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=766, - serialized_end=807, + serialized_start=749, + serialized_end=790, ) @@ -416,8 +409,8 @@ _ADDRESS = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=809, - serialized_end=846, + serialized_start=792, + serialized_end=829, ) @@ -531,8 +524,8 @@ _CLIENTCONNECTION = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=849, - serialized_end=1171, + serialized_start=832, + serialized_end=1154, ) @@ -653,8 +646,8 @@ _SERVERCONNECTION = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=1174, - serialized_end=1537, + serialized_start=1157, + serialized_end=1520, ) @@ -691,8 +684,8 @@ _TLSEXTENSION = _descriptor.Descriptor( extension_ranges=[], oneofs=[ ], - serialized_start=1539, - serialized_end=1581, + serialized_start=1522, + serialized_end=1564, ) _HTTPFLOW.fields_by_name['request'].message_type = _HTTPREQUEST diff --git a/mitmproxy/io/protobuf.py b/mitmproxy/io/protobuf.py index 3db9a013..9a00eacf 100644 --- a/mitmproxy/io/protobuf.py +++ b/mitmproxy/io/protobuf.py @@ -2,7 +2,6 @@ import typing from mitmproxy import flow from mitmproxy import exceptions -from mitmproxy import ctx from mitmproxy.http import HTTPFlow, HTTPResponse, HTTPRequest from mitmproxy.certs import Cert from mitmproxy.connections import ClientConnection, ServerConnection @@ -16,7 +15,8 @@ def _move_attrs(s_obj, d_obj, attrs): setattr(d_obj, attr, getattr(s_obj, attr)) else: if hasattr(s_obj, attr) and getattr(s_obj, attr) is not None: - if not getattr(s_obj, attr): + # ugly fix to set None in empty str or bytes fields + if getattr(s_obj, attr) == "" or getattr(s_obj, attr) == b"": d_obj[attr] = None else: d_obj[attr] = getattr(s_obj, attr) @@ -87,12 +87,12 @@ def _dump_http_error(e: flow.Error) -> http_pb2.HTTPError: return pe -def dump_http(f: HTTPFlow) -> http_pb2.HTTPFlow(): +def dump_http(f: flow.Flow) -> http_pb2.HTTPFlow: pf = http_pb2.HTTPFlow() for p in ['request', 'response', 'client_conn', 'server_conn', 'error']: - if hasattr(f, p): + if hasattr(f, p) and getattr(f, p): getattr(pf, p).MergeFrom(eval(f"_dump_http_{p}")(getattr(f, p))) - _move_attrs(f, pf, ['intercepted', 'marked', 'mode', 'id', 'version']) + _move_attrs(f, pf, ['intercepted', 'marked', 'mode', 'id']) return pf @@ -105,9 +105,9 @@ def dumps(f: flow.Flow) -> bytes: def _load_http_request(o: http_pb2.HTTPRequest) -> HTTPRequest: - d = {} + d: dict = {} _move_attrs(o, d, ['first_line_format', 'method', 'scheme', 'host', 'port', 'path', 'http_version', 'content', - 'timestamp_start', 'timestamp_end', 'is_replay']) + 'timestamp_start', 'timestamp_end', 'is_replay']) if d['content'] is None: d['content'] = b"" d["headers"] = [] @@ -118,9 +118,9 @@ def _load_http_request(o: http_pb2.HTTPRequest) -> HTTPRequest: def _load_http_response(o: http_pb2.HTTPResponse) -> HTTPResponse: - d = {} + d: dict = {} _move_attrs(o, d, ['http_version', 'status_code', 'reason', - 'content', 'timestamp_start', 'timestamp_end', 'is_replay']) + 'content', 'timestamp_start', 'timestamp_end', 'is_replay']) if d['content'] is None: d['content'] = b"" d["headers"] = [] @@ -131,14 +131,12 @@ def _load_http_response(o: http_pb2.HTTPResponse) -> HTTPResponse: def _load_http_client_conn(o: http_pb2.ClientConnection) -> ClientConnection: - d = {} - _move_attrs(o, d, ['id', 'tls_established', 'sni', 'alpn_proto_negotiated', 'tls_version', - 'timestamp_start', 'timestamp_tcp_setup', 'timestamp_tls_setup', 'timestamp_end']) + d: dict = {} + _move_attrs(o, d, ['id', 'tls_established', 'sni', 'cipher_name', 'alpn_proto_negotiated', 'tls_version', + 'timestamp_start', 'timestamp_tcp_setup', 'timestamp_tls_setup', 'timestamp_end']) for cert in ['clientcert', 'mitmcert']: if hasattr(o, cert) and getattr(o, cert): - c = Cert("") - c.from_pem(getattr(o, cert)) - d[cert] = c + d[cert] = Cert.from_pem(getattr(o, cert)) if o.tls_extensions: d['tls_extensions'] = [] for extension in o.tls_extensions: @@ -152,18 +150,17 @@ def _load_http_client_conn(o: http_pb2.ClientConnection) -> ClientConnection: def _load_http_server_conn(o: http_pb2.ServerConnection) -> ServerConnection: - d = {} + d: dict = {} _move_attrs(o, d, ['id', 'tls_established', 'sni', 'alpn_proto_negotiated', 'tls_version', - 'timestamp_start', 'timestamp_tcp_setup', 'timestamp_tls_setup', 'timestamp_end']) + 'timestamp_start', 'timestamp_tcp_setup', 'timestamp_tls_setup', 'timestamp_end']) for addr in ['address', 'ip_address', 'source_address']: if hasattr(o, addr): d[addr] = (getattr(o, addr).host, getattr(o, addr).port) if o.cert: - c = Cert("") - c.from_pem(o.cert) + c = Cert.from_pem(o.cert) d['cert'] = c - if len(o.via.id): - d['via'] = _load_http_server_conn(d['via']) + if o.HasField('via'): + d['via'] = _load_http_server_conn(o.via) sc = ServerConnection(tuple()) for k, v in d.items(): setattr(sc, k, v) @@ -181,9 +178,11 @@ def _load_http_error(o: http_pb2.HTTPError) -> typing.Optional[flow.Error]: def load_http(hf: http_pb2.HTTPFlow) -> HTTPFlow: parts = {} for p in ['request', 'response', 'client_conn', 'server_conn', 'error']: - if hasattr(hf, p) and getattr(hf, p): + if hf.HasField(p): parts[p] = eval(f"_load_http_{p}")(getattr(hf, p)) - _move_attrs(hf, parts, ['intercepted', 'marked', 'mode', 'id', 'version']) + else: + parts[p] = None + _move_attrs(hf, parts, ['intercepted', 'marked', 'mode', 'id']) f = HTTPFlow(ClientConnection(None, tuple(), None), ServerConnection(tuple())) for k, v in parts.items(): setattr(f, k, v) @@ -195,9 +194,5 @@ def loads(b: bytes, typ="http") -> flow.Flow: raise exceptions.TypeError("Flow types different than HTTP not supported yet!") else: p = http_pb2.HTTPFlow() - try: - p.ParseFromString(b) - return load_http(p) - except Exception as e: - ctx.log(str(e)) - + p.ParseFromString(b) + return load_http(p) @@ -2,7 +2,7 @@ max-line-length = 140 max-complexity = 25 ignore = E251,C901,W503,W292,E722,E741 -exclude = mitmproxy/contrib/*,test/mitmproxy/data/*,release/build/* +exclude = mitmproxy/contrib/*,test/mitmproxy/data/*,release/build/*,mitmproxy/io/proto/* addons = file,open,basestring,xrange,unicode,long,cmp [tool:pytest] @@ -70,11 +70,11 @@ setup( "kaitaistruct>=0.7,<0.9", "ldap3>=2.5,<2.6", "passlib>=1.6.5, <1.8", + "protobuf>=3.6.0, <3.7", "pyasn1>=0.3.1,<0.5", "pyOpenSSL>=17.5,<18.1", "pyparsing>=2.1.3, <2.3", "pyperclip>=1.6.0, <1.7", - "protobuf>=3.6.0, <3.7", "ruamel.yaml>=0.13.2, <0.16", "sortedcontainers>=1.5.4,<2.1", "tornado>=4.3,<5.1", diff --git a/test/filename_matching.py b/test/filename_matching.py index 5f49725e..f5321307 100755 --- a/test/filename_matching.py +++ b/test/filename_matching.py @@ -9,7 +9,8 @@ import sys def check_src_files_have_test(): missing_test_files = [] - excluded = ['mitmproxy/contrib/', 'mitmproxy/test/', 'mitmproxy/tools/', 'mitmproxy/platform/'] + excluded = ['mitmproxy/contrib/', 'mitmproxy/io/proto/', + 'mitmproxy/test/', 'mitmproxy/tools/', 'mitmproxy/platform/'] src_files = glob.glob('mitmproxy/**/*.py', recursive=True) + glob.glob('pathod/**/*.py', recursive=True) src_files = [f for f in src_files if os.path.basename(f) != '__init__.py'] src_files = [f for f in src_files if not any(os.path.normpath(p) in f for p in excluded)] diff --git a/test/mitmproxy/io/test_db.py b/test/mitmproxy/io/test_db.py new file mode 100644 index 00000000..4a2dfb67 --- /dev/null +++ b/test/mitmproxy/io/test_db.py @@ -0,0 +1,26 @@ +from mitmproxy.io import db +from mitmproxy.test import tflow + + +class TestDB: + + def test_create(self, tdata): + dh = db.DBHandler(db_path=tdata.path("mitmproxy/data") + "/tmp.sqlite") + with dh._con as c: + cur = c.cursor() + cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='FLOWS';") + assert cur.fetchall() == [('FLOWS',)] + + def test_roundtrip(self, tdata): + dh = db.DBHandler(db_path=tdata.path("mitmproxy/data") + "/tmp.sqlite", mode='write') + flows = [] + for i in range(10): + flows.append(tflow.tflow()) + dh.store(flows) + dh = db.DBHandler(db_path=tdata.path("mitmproxy/data") + "/tmp.sqlite") + with dh._con as c: + cur = c.cursor() + cur.execute("SELECT count(*) FROM FLOWS;") + assert cur.fetchall()[0][0] == 10 + loaded_flows = dh.load() + assert len(loaded_flows) == len(flows) diff --git a/test/mitmproxy/io/test_protobuf.py b/test/mitmproxy/io/test_protobuf.py new file mode 100644 index 00000000..f725b980 --- /dev/null +++ b/test/mitmproxy/io/test_protobuf.py @@ -0,0 +1,120 @@ +import pytest + +from mitmproxy import certs +from mitmproxy import http +from mitmproxy import exceptions +from mitmproxy.test import tflow, tutils +from mitmproxy.io import protobuf + + +class TestProtobuf: + + def test_roundtrip_client(self): + c = tflow.tclient_conn() + del c.reply + c.rfile = None + c.wfile = None + pc = protobuf._dump_http_client_conn(c) + lc = protobuf._load_http_client_conn(pc) + assert c.__dict__ == lc.__dict__ + + def test_roundtrip_client_cert(self, tdata): + c = tflow.tclient_conn() + c.rfile = None + c.wfile = None + del c.reply + with open(tdata.path("mitmproxy/net/data/clientcert/client.pem"), "rb") as f: + d = f.read() + c.clientcert = certs.Cert.from_pem(d) + pc = protobuf._dump_http_client_conn(c) + lc = protobuf._load_http_client_conn(pc) + assert c.__dict__ == lc.__dict__ + + def test_roundtrip_server(self): + s = tflow.tserver_conn() + del s.reply + s.wfile = None + s.rfile = None + ps = protobuf._dump_http_server_conn(s) + ls = protobuf._load_http_server_conn(ps) + assert s.__dict__ == ls.__dict__ + + def test_roundtrip_server_cert(self, tdata): + s = tflow.tserver_conn() + del s.reply + s.wfile = None + s.rfile = None + with open(tdata.path("mitmproxy/net/data/text_cert"), "rb") as f: + d = f.read() + s.cert = certs.Cert.from_pem(d) + ps = protobuf._dump_http_server_conn(s) + ls = protobuf._load_http_server_conn(ps) + assert s.__dict__ == ls.__dict__ + + def test_roundtrip_server_via(self): + s = tflow.tserver_conn() + s.via = tflow.tserver_conn() + del s.reply + s.wfile = None + s.rfile = None + ps = protobuf._dump_http_server_conn(s) + ls = protobuf._load_http_server_conn(ps) + assert s.__dict__ == ls.__dict__ + del s.via.reply + s.via.wfile = None + s.via.rfile = None + assert s.via.__dict__ == ls.via.__dict__ + + def test_roundtrip_http_request(self): + req = http.HTTPRequest.wrap(tutils.treq()) + preq = protobuf._dump_http_request(req) + lreq = protobuf._load_http_request(preq) + assert req.__dict__ == lreq.__dict__ + + def test_roundtrip_http_request_empty_content(self): + req = http.HTTPRequest.wrap(tutils.treq(content=b"")) + preq = protobuf._dump_http_request(req) + lreq = protobuf._load_http_request(preq) + assert req.__dict__ == lreq.__dict__ + + def test_roundtrip_http_response(self): + res = http.HTTPResponse.wrap(tutils.tresp()) + pres = protobuf._dump_http_response(res) + lres = protobuf._load_http_response(pres) + assert res.__dict__ == lres.__dict__ + + def test_roundtrip_http_response_empty_content(self): + res = http.HTTPResponse.wrap(tutils.tresp(content=b"")) + pres = protobuf._dump_http_response(res) + lres = protobuf._load_http_response(pres) + assert res.__dict__ == lres.__dict__ + + def test_roundtrip_http_error(self): + err = tflow.terr() + perr = protobuf._dump_http_error(err) + lerr = protobuf._load_http_error(perr) + assert err.__dict__ == lerr.__dict__ + + def test_roundtrip_http_flow_only_req(self): + f = tflow.tflow() + f.reply = None + pf = protobuf.dumps(f) + lf = protobuf.loads(pf, "http") + assert f.__dict__ == lf.__dict__ + + def test_roundtrip_http_flow_res(self): + f = tflow.tflow(resp=True) + f.reply = None + pf = protobuf.dumps(f) + lf = protobuf.loads(pf, "http") + assert f.__dict__ == lf.__dict__ + + def test_unsupported_dumps(self): + w = tflow.twebsocketflow() + with pytest.raises(exceptions.TypeError): + protobuf.dumps(w) + + def test_unsupported_loads(self): + b = b"blobs" + with pytest.raises(exceptions.TypeError): + protobuf.loads(b, 'not-http') |