simplify server changes for inline scripts

1 files changed, 27 insertions, 24 deletions

diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index a84b4061..ef5c87fb 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -2,7 +2,6 @@ from __future__ import absolute_import
 import copy
 import netlib.tcp
 from .. import stateobject, utils, version
-from ..proxy.primitives import AddressPriority
 from ..proxy.connection import ClientConnection, ServerConnection
 
 
@@ -153,44 +152,48 @@ class LiveConnection(object):
     without requiring the expose the ConnectionHandler.
     """
     def __init__(self, c):
-        self._c = c
+        self.c = c
+        self._backup_server_conn = None
         """@type: libmproxy.proxy.server.ConnectionHandler"""
 
-    def change_server(self, address, ssl, persistent_change=False):
+    def change_server(self, address, ssl=False, force=False, persistent_change=False):
         address = netlib.tcp.Address.wrap(address)
-        if address != self._c.server_conn.address:
+        if force or address != self.c.server_conn.address or ssl != self.c.server_conn.ssl_established:
 
-            self._c.log("Change server connection: %s:%s -> %s:%s" % (
-                self._c.server_conn.address.host,
-                self._c.server_conn.address.port,
+            self.c.log("Change server connection: %s:%s -> %s:%s [persistent: %s]" % (
+                self.c.server_conn.address.host,
+                self.c.server_conn.address.port,
                 address.host,
-                address.port
+                address.port,
+                persistent_change
             ), "debug")
 
-            if not hasattr(self, "_backup_server_conn"):
-                self._backup_server_conn = self._c.server_conn
-                self._c.server_conn = None
+            if self._backup_server_conn:
+                self._backup_server_conn = self.c.server_conn
+                self.c.server_conn = None
             else:  # This is at least the second temporary change. We can kill the current connection.
-                self._c.del_server_connection()
+                self.c.del_server_connection()
 
-            self._c.set_server_address(address, AddressPriority.MANUALLY_CHANGED)
-            self._c.establish_server_connection(ask=False)
+            self.c.set_server_address(address)
+            self.c.establish_server_connection(ask=False)
             if ssl:
-                self._c.establish_ssl(server=True)
-        if hasattr(self, "_backup_server_conn") and persistent_change:
-            del self._backup_server_conn
+                self.c.establish_ssl(server=True)
+        if persistent_change:
+            self._backup_server_conn = None
 
     def restore_server(self):
-        if not hasattr(self, "_backup_server_conn"):
+        # TODO: Similar to _backup_server_conn, introduce _cache_server_conn, which keeps the changed connection open
+        # This may be beneficial if a user is rewriting all requests from http to https or similar.
+        if not self._backup_server_conn:
             return
 
-        self._c.log("Restore original server connection: %s:%s -> %s:%s" % (
-            self._c.server_conn.address.host,
-            self._c.server_conn.address.port,
+        self.c.log("Restore original server connection: %s:%s -> %s:%s" % (
+            self.c.server_conn.address.host,
+            self.c.server_conn.address.port,
             self._backup_server_conn.address.host,
             self._backup_server_conn.address.port
         ), "debug")
 
-        self._c.del_server_connection()
-        self._c.server_conn = self._backup_server_conn
-        del self._backup_server_conn
\ No newline at end of file
+        self.c.del_server_connection()
+        self.c.server_conn = self._backup_server_conn
+        self._backup_server_conn = None
\ No newline at end of file