Format.

src/zamenyat/__main__.py

@@ -1,35 +1,22 @@
 import argparse
+
 from zamenyat.app import Application
 
 parser = argparse.ArgumentParser(description="Zamenyat sensitive content replacer.")
 
-parser.add_argument(
-    "--host",
-    required=True,
-    type=str
-)
-parser.add_argument(
-    "--port",
-    required=True,
-    type=int
-)
-parser.add_argument(
-    "--upstream-host",
-    required=True,
-    type=str
-)
-parser.add_argument(
-    "--upstream-port",
-    required=True,
-    type=int
-)
+parser.add_argument("--host", required=True, type=str)
+parser.add_argument("--port", required=True, type=int)
+parser.add_argument("--upstream-host", required=True, type=str)
+parser.add_argument("--upstream-port", required=True, type=int)
 
 
 def main():
     args = parser.parse_args()
-    app = Application(upstream_host=args.upstream_host, upstream_port=args.upstream_port)
+    app = Application(
+        upstream_host=args.upstream_host, upstream_port=args.upstream_port
+    )
     app.serve(host=args.host, port=args.port)
 
-if __name__ == '__main__':
-    main()
 
+if __name__ == "__main__":
+    main()

src/zamenyat/app.py

@@ -1,13 +1,15 @@
-from app.app import Application as BaseApplication, get_timestamp
 import asyncio
-from concurrent.futures import ThreadPoolExecutor as Executor
 import time
+from concurrent.futures import ThreadPoolExecutor as Executor
+
+from app.app import get_timestamp
 
 ZAMENYAT_BACKLOG = 100
 ZAMENYAT_THREAD_COUNT = 2500
 ZAMENYAT_BUFFER_SIZE = 4096
 ZAMENYAT_HEADER_MAX_LENGTH = 4096 * 2
 
+
 class AsyncWriter:
 
     def __init__(self, writer, buffer_size=ZAMENYAT_BUFFER_SIZE, debug=False):
@@ -29,16 +31,17 @@ class AsyncWriter:
             data = data[chunk_size:]
             await self.writer.drain()
 
+
 class AsyncReader:
 
     def __init__(self, reader, debug=False):
         self.reader = reader
-        self.buffer = b''
+        self.buffer = b""
         self.debug = debug
 
     async def read_until(self, to_match):
-        buffer = b''
-        while not to_match in buffer:
+        buffer = b""
+        while to_match not in buffer:
             chunk = await self.read()
             if not chunk:
                 return None
@@ -49,7 +52,7 @@ class AsyncReader:
         return data
 
     async def read(self, buffer_size=ZAMENYAT_BUFFER_SIZE, exact=False):
-        read_extra = buffer_size - len(self.buffer)
+        buffer_size - len(self.buffer)
         while len(self.buffer) < buffer_size:
             chunk_size = buffer_size - len(self.buffer)
             chunk = await self.reader.read(chunk_size)
@@ -61,7 +64,9 @@ class AsyncReader:
             self.buffer += chunk
             if not exact:
                 break
-        buffer_size = len(self.buffer) if len(self.buffer) < buffer_size else buffer_size
+        buffer_size = (
+            len(self.buffer) if len(self.buffer) < buffer_size else buffer_size
+        )
         data = self.buffer[:buffer_size]
         self.buffer = self.buffer[buffer_size:]
 
@@ -70,10 +75,11 @@ class AsyncReader:
     async def unread(self, data):
         if not data:
             return
-        if hasattr(data, 'encode'):
+        if hasattr(data, "encode"):
             data = data.encode()
         self.buffer = data + self.buffer
 
+
 class Socket:
 
     def __init__(self, reader, writer, buffer_size, debug=False):
@@ -88,12 +94,14 @@ class Socket:
         self.close = self.writer.close
         self.wait_closed = self.writer.wait_closed
 
+
 class Application:
 
-    def __init__(self, upstream_host, upstream_port, *args, **kwargs):
+    def __init__(self, upstream_host, upstream_port, silent=False, *args, **kwargs):
         self.upstream_host = upstream_host
         self.upstream_port = upstream_port
         self.server = None
+        self.silent = silent
         self.host = None
         self.port = None
         self.executor = None
@@ -104,8 +112,7 @@ class Application:
         super().__init__(*args, **kwargs)
 
     async def get_headers(self, reader):
-        data = b''
-        headers = await reader.read_until(b'\r\n\r\n')
+        headers = await reader.read_until(b"\r\n\r\n")
         if not headers:
             return None, None
         headers = headers[:-2]
@@ -121,7 +128,7 @@ class Application:
     def header_dict_to_bytes(self, req_resp, headers):
         header_list = [req_resp]
         for key, value in headers.items():
-            header_list.append("{}: {}".format(key, value))
+            header_list.append(f"{key}: {value}")
         header_list.append("\r\n")
         return ("\r\n".join(header_list)).encode()
 
@@ -132,15 +139,20 @@ class Application:
             writer = Socket(reader, writer, ZAMENYAT_BUFFER_SIZE)
             while True:
                 req_resp, headers = None, None
-                data = b''
+                data = b""
                 if not is_websocket:
                     req_resp, headers = await self.get_headers(reader)
                     if not headers:
                         return None
                 if headers:
-                    if 'Content-Length' in headers:
-                        while len(data) != headers['Content-Length']:
-                            chunk_size = headers['Content-Length'] - len(data) if self.buffer_size > headers['Content-Length'] - len(data) else self.buffer_size
+                    if "Content-Length" in headers:
+                        while len(data) != headers["Content-Length"]:
+                            chunk_size = (
+                                headers["Content-Length"] - len(data)
+                                if self.buffer_size
+                                > headers["Content-Length"] - len(data)
+                                else self.buffer_size
+                            )
                             chunk = await reader.read(chunk_size)
                             if not chunk:
                                 data = None
@@ -168,36 +180,50 @@ class Application:
         self.total_connection_count += 1
         connection_nr = self.total_connection_count
 
-
-        upstream_reader, upstream_writer = await asyncio.open_connection(self.upstream_host, self.upstream_port)
+        upstream_reader, upstream_writer = await asyncio.open_connection(
+            self.upstream_host, self.upstream_port
+        )
 
         is_websocket = False
 
         while True:
             time_start = time.time()
-            print(f"Connected to upstream #{connection_nr} server {self.upstream_host}:{self.upstream_port} #{self.connection_count} Time: {get_timestamp()}")
+            if not self.silent:
+                print(
+                    f"Connected to upstream #{connection_nr} server {self.upstream_host}:{self.upstream_port} #{self.connection_count} Time: {get_timestamp()}"
+                )
 
             if is_websocket:
                 await asyncio.gather(
                     self.stream(reader, upstream_writer, is_websocket),
-                    self.stream(upstream_reader, writer, is_websocket)
+                    self.stream(upstream_reader, writer, is_websocket),
+                )
+                if not self.silent:
+                    print(
+                        f"Disconnected upstream #{connection_nr} server {self.upstream_host}:{self.upstream_port} #{self.connection_count} Duration: {time_duration:.5f}s"
                     )
                 break
             else:
-                request_headers = await self.stream(reader, upstream_writer,is_websocket) 
+                request_headers = await self.stream(
+                    reader, upstream_writer, is_websocket
+                )
                 await self.stream(upstream_reader, writer, is_websocket)
 
                 keep_alive = False
                 if request_headers:
-                    if request_headers.get('Connection') == 'keep-alive': # and not headers.get('Upgrade-Insecure-Requests'):
+                    if (
+                        request_headers.get("Connection") == "keep-alive"
+                    ):  # and not headers.get('Upgrade-Insecure-Requests'):
                         keep_alive = True
-                    if request_headers.get("Upgrade") == 'websocket':
+                    if request_headers.get("Upgrade") == "websocket":
                         is_websocket = True
 
                 time_end = time.time()
                 time_duration = time_end - time_start
-                print(f"Disconnected upstream #{connection_nr} server {self.upstream_host}:{self.upstream_port} #{self.connection_count} Duration: {time_duration:.5f}s")
-                
+                if not self.silent:
+                    print(
+                        f"Disconnected upstream #{connection_nr} server {self.upstream_host}:{self.upstream_port} #{self.connection_count} Duration: {time_duration:.5f}s"
+                    )
 
                 if not any([keep_alive, is_websocket]):
                     break
@@ -209,9 +235,6 @@ class Application:
         await writer.wait_closed()
         await upstream_writer.wait_closed()
 
-
-
-
     def upgrade_executor(self, thread_count):
         self.executor = Executor(max_workers=thread_count)
         loop = asyncio.get_running_loop()
@@ -222,7 +245,9 @@ class Application:
         self.upgrade_executor(ZAMENYAT_THREAD_COUNT)
         self.host = host
         self.port = port
-        self.server = await asyncio.start_server(self.handle_client, self.host, self.port,backlog=backlog)
+        self.server = await asyncio.start_server(
+            self.handle_client, self.host, self.port, backlog=backlog
+        )
         async with self.server:
             await self.server.serve_forever()