From 47651b1ff2e5dd0c788d744dc1b1d5587107f71a Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Thu, 28 Jun 2012 14:29:15 +1200
Subject: [PATCH] Serialization and de-serialization of new cert format.

---
 libmproxy/flow.py | 17 +++++------------
 test/test_flow.py | 12 +-----------
 test/tutils.py    |  4 +++-
 3 files changed, 9 insertions(+), 24 deletions(-)

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 0e5ba0e36..4de8bc96e 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -463,12 +463,12 @@ class Response(HTTPMsg):
 
             timestamp: Seconds since the epoch
     """
-    def __init__(self, request, httpversion, code, msg, headers, content, der_cert, timestamp=None):
+    def __init__(self, request, httpversion, code, msg, headers, content, cert, timestamp=None):
         assert isinstance(headers, ODictCaseless)
         self.request = request
         self.httpversion, self.code, self.msg = httpversion, code, msg
         self.headers, self.content = headers, content
-        self.der_cert = der_cert
+        self.cert = cert
         self.timestamp = timestamp or utils.timestamp()
         controller.Msg.__init__(self)
         self.replay = False
@@ -538,14 +538,7 @@ class Response(HTTPMsg):
         self.headers = ODictCaseless._from_state(state["headers"])
         self.content = state["content"]
         self.timestamp = state["timestamp"]
-        self.der_cert = state["der_cert"]
-
-    def get_cert(self):
-        """
-            Returns a certutils.SSLCert object, or None.
-        """
-        if self.der_cert:
-            return certutils.SSLCert.from_der(self.der_cert)
+        self.cert = state["cert"]
 
     def _get_state(self):
         return dict(
@@ -554,7 +547,7 @@ class Response(HTTPMsg):
             msg = self.msg,
             headers = self.headers._get_state(),
             timestamp = self.timestamp,
-            der_cert = self.der_cert,
+            cert = self.cert.to_pem() if self.cert else None,
             content = self.content
         )
 
@@ -567,7 +560,7 @@ class Response(HTTPMsg):
             str(state["msg"]),
             ODictCaseless._from_state(state["headers"]),
             state["content"],
-            state.get("der_cert"),
+            certutils.SSLCert.from_pem(state["cert"]) if state["cert"] else None,
             state["timestamp"],
         )
 
diff --git a/test/test_flow.py b/test/test_flow.py
index 89577ddbc..1394a6685 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -2,6 +2,7 @@ import Queue, time
 from cStringIO import StringIO
 import email.utils
 from libmproxy import filt, flow, controller, utils, tnetstring
+from netlib import certutils
 import tutils
 import libpry
 
@@ -865,17 +866,6 @@ class TestResponse:
         c = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
         assert "00:21:38" in r._refresh_cookie(c, 60)
 
-    def test_get_cert(self):
-        req = tutils.treq()
-        resp = flow.Response(
-            req, (1, 1), 200, "msg", flow.ODictCaseless(), "content",
-            file(tutils.test_data.path("data/dercert")).read()
-        )
-        assert resp.get_cert()
-
-        resp = tutils.tresp()
-        assert not resp.get_cert()
-
     def test_getset_state(self):
         h = flow.ODictCaseless()
         h["test"] = ["test"]
diff --git a/test/tutils.py b/test/tutils.py
index da1cf6625..79a1da39f 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -2,6 +2,7 @@ import threading, Queue, time
 import os, shutil, tempfile
 from contextlib import contextmanager
 from libmproxy import proxy, flow, controller, utils
+from netlib import certutils
 import human_curl as hurl
 import libpathod.test
 import random
@@ -19,7 +20,8 @@ def tresp(req=None):
         req = treq()
     headers = flow.ODictCaseless()
     headers["header_response"] = ["svalue"]
-    return flow.Response(req, (1, 1), 200, "message", headers, "content_response", None)
+    cert = certutils.SSLCert.from_der(file(test_data.path("data/dercert")).read())
+    return flow.Response(req, (1, 1), 200, "message", headers, "content_response", cert)
 
 
 def tflow():