protobuf: tests implemented, full coverage

5 files changed, 49 insertions, 61 deletions

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)