Code factorization

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.4.1
+current_version = 0.4.2
 commit = True
 tag = True
 

CHANGELOG.md

@@ -1,3 +1,9 @@
+## [0.4.1] - 2025-06-05
+### Added
+- Custom log format in config.ini
+- Add __main__.py to run module
+- Colored log
+
 ## [0.3.3] - 2025-05-30
 ### Added
 - Change pylint to flake8

pyproxy/__init__.py

@@ -5,7 +5,7 @@
 
 import os
 
-__version__ = "0.4.1"
+__version__ = "0.4.2"
 
 if os.path.isdir("pyproxy/monitoring"):
     __slim__ = False

pyproxy/handlers/client.py

@@ -39,9 +39,7 @@ def __init__(
         shortcuts,
         custom_header,
         active_connections,
-        proxy_enable,
-        proxy_host,
-        proxy_port,
+        proxy_config,
     ):
         self.html_403 = html_403
         self.logger_config = logger_config
@@ -58,11 +56,32 @@ def __init__(
         self.console_logger = console_logger
         self.config_shortcuts = shortcuts
         self.config_custom_header = custom_header
-        self.proxy_enable = proxy_enable
-        self.proxy_host = proxy_host
-        self.proxy_port = proxy_port
+        self.proxy_config = proxy_config
         self.active_connections = active_connections
 
+    def _create_handler(self, handler_class, **extra_kwargs):
+        """
+        Factory to create handler instance with shared common parameters.
+        """
+        params = dict(
+            html_403=self.html_403,
+            logger_config=self.logger_config,
+            filter_config=self.filter_config,
+            filter_queue=self.filter_queue,
+            filter_result_queue=self.filter_result_queue,
+            shortcuts_queue=self.shortcuts_queue,
+            shortcuts_result_queue=self.shortcuts_result_queue,
+            custom_header_queue=self.custom_header_queue,
+            custom_header_result_queue=self.custom_header_result_queue,
+            console_logger=self.console_logger,
+            shortcuts=self.config_shortcuts,
+            custom_header=self.config_custom_header,
+            proxy_config=self.proxy_config,
+            active_connections=self.active_connections,
+        )
+        params.update(extra_kwargs)
+        return handler_class(**params)
+
     def handle_client(self, client_socket):
         """
         Handles an incoming client connection by processing the request and forwarding
@@ -82,45 +101,13 @@ def handle_client(self, client_socket):
         first_line = request.decode(errors="ignore").split("\n")[0]
 
         if first_line.startswith("CONNECT"):
-            client_https_handler = HttpsHandler(
-                html_403=self.html_403,
-                logger_config=self.logger_config,
-                filter_config=self.filter_config,
+            https_handler = self._create_handler(
+                HttpsHandler,
                 ssl_config=self.ssl_config,
-                filter_queue=self.filter_queue,
-                filter_result_queue=self.filter_result_queue,
-                shortcuts_queue=self.shortcuts_queue,
-                shortcuts_result_queue=self.shortcuts_result_queue,
                 cancel_inspect_queue=self.cancel_inspect_queue,
                 cancel_inspect_result_queue=self.cancel_inspect_result_queue,
-                custom_header_queue=self.custom_header_queue,
-                custom_header_result_queue=self.custom_header_result_queue,
-                console_logger=self.console_logger,
-                shortcuts=self.config_shortcuts,
-                custom_header=self.config_custom_header,
-                proxy_enable=self.proxy_enable,
-                proxy_host=self.proxy_host,
-                proxy_port=self.proxy_port,
-                active_connections=self.active_connections,
             )
-            client_https_handler.handle_https_connection(client_socket, first_line)
+            https_handler.handle_https_connection(client_socket, first_line)
         else:
-            client_http_handler = HttpHandler(
-                html_403=self.html_403,
-                logger_config=self.logger_config,
-                filter_config=self.filter_config,
-                filter_queue=self.filter_queue,
-                filter_result_queue=self.filter_result_queue,
-                shortcuts_queue=self.shortcuts_queue,
-                shortcuts_result_queue=self.shortcuts_result_queue,
-                custom_header_queue=self.custom_header_queue,
-                custom_header_result_queue=self.custom_header_result_queue,
-                console_logger=self.console_logger,
-                shortcuts=self.config_shortcuts,
-                custom_header=self.config_custom_header,
-                proxy_enable=self.proxy_enable,
-                proxy_host=self.proxy_host,
-                proxy_port=self.proxy_port,
-                active_connections=self.active_connections,
-            )
-            client_http_handler.handle_http_request(client_socket, request)
+            http_handler = self._create_handler(HttpHandler)
+            http_handler.handle_http_request(client_socket, request)

pyproxy/handlers/http.py

@@ -8,8 +8,9 @@
 import socket
 import os
 import threading
+from urllib.parse import urlparse
 
-from pyproxy.utils.http_req import extract_headers, parse_url
+from pyproxy.utils.http_req import extract_headers
 
 
 class HttpHandler:
@@ -34,9 +35,7 @@ def __init__(
         shortcuts,
         custom_header,
         active_connections,
-        proxy_enable,
-        proxy_host,
-        proxy_port,
+        proxy_config,
     ):
         self.html_403 = html_403
         self.logger_config = logger_config
@@ -50,11 +49,81 @@ def __init__(
         self.console_logger = console_logger
         self.config_shortcuts = shortcuts
         self.config_custom_header = custom_header
-        self.proxy_enable = proxy_enable
-        self.proxy_host = proxy_host
-        self.proxy_port = proxy_port
+        self.proxy_config = proxy_config
         self.active_connections = active_connections
 
+    def _get_modified_headers(self, url, request_text):
+        """
+        Extract headers from a request
+        """
+        headers = extract_headers(request_text)
+        self.custom_header_queue.put(url)
+        try:
+            new_headers = self.custom_header_result_queue.get(timeout=5)
+            headers.update(new_headers)
+        except Exception:
+            self.console_logger.warning(
+                "Timeout while getting custom headers for %s", url
+            )
+        return headers
+
+    def _rebuild_http_request(self, request_line, headers, body=""):
+        """
+        Reconstructs an HTTP request with the new headers.
+        """
+        header_lines = [f"{key}: {value}" for key, value in headers.items()]
+        reconstructed_headers = "\r\n".join(header_lines)
+        return f"{request_line}\r\n{reconstructed_headers}\r\n\r\n{body}".encode()
+
+    def _apply_shortcut(self, url: str) -> str | None:
+        """
+        Checks if a shortcut is defined for the given domain.
+        """
+        if self.config_shortcuts and os.path.isfile(self.config_shortcuts):
+            parsed_url = urlparse(url)
+            domain = parsed_url.hostname
+            self.shortcuts_queue.put(domain)
+            try:
+                return self.shortcuts_result_queue.get(timeout=5)
+            except Exception:
+                self.console_logger.warning(
+                    "Timeout while getting shortcut for %s", url
+                )
+        return None
+
+    def _is_blocked(self, url: str) -> bool:
+        """
+        Checks if a URL is blocked by the configuration filter.
+        """
+        if not self.filter_config.no_filter:
+            self.filter_queue.put(url)
+            try:
+                result = self.filter_result_queue.get(timeout=5)
+                return result[1] == "Blocked"
+            except Exception:
+                self.console_logger.warning("Timeout while filtering %s", url)
+        return False
+
+    def _send_403(self, client_socket, url, first_line):
+        """
+        Sends an HTTP 403 Forbidden response to the client.
+        """
+        if not self.logger_config.no_logging_block:
+            self.logger_config.block_logger.info(
+                "%s - %s - %s", client_socket.getpeername()[0], url, first_line
+            )
+        with open(self.html_403, "r", encoding="utf-8") as f:
+            custom_403_page = f.read()
+        response = (
+            f"HTTP/1.1 403 Forbidden\r\n"
+            f"Content-Length: {len(custom_403_page)}\r\n"
+            f"\r\n"
+            f"{custom_403_page}"
+        )
+        client_socket.sendall(response.encode())
+        client_socket.close()
+        self.active_connections.pop(threading.get_ident(), None)
+
     def handle_http_request(self, client_socket, request):
         """
         Processes an HTTP request, checks for URL filtering, applies shortcuts,
@@ -67,16 +136,8 @@ def handle_http_request(self, client_socket, request):
         first_line = request.decode(errors="ignore").split("\n")[0]
         url = first_line.split(" ")[1]
 
-        if self.config_custom_header and os.path.isfile(self.config_custom_header):
-            headers = extract_headers(request.decode(errors="ignore"))
-            self.custom_header_queue.put(url)
-            new_headers = self.custom_header_result_queue.get(timeout=5)
-            headers.update(new_headers)
-
         if self.config_shortcuts and os.path.isfile(self.config_shortcuts):
-            domain, _ = parse_url(url)
-            self.shortcuts_queue.put(domain)
-            shortcut_url = self.shortcuts_result_queue.get(timeout=5)
+            shortcut_url = self._apply_shortcut(url)
             if shortcut_url:
                 response = (
                     f"HTTP/1.1 302 Found\r\n"
@@ -90,27 +151,12 @@ def handle_http_request(self, client_socket, request):
                 self.active_connections.pop(threading.get_ident(), None)
                 return
 
-        if not self.filter_config.no_filter:
-            self.filter_queue.put(url)
-            result = self.filter_result_queue.get(timeout=5)
-            if result[1] == "Blocked":
-                if not self.logger_config.no_logging_block:
-                    self.logger_config.block_logger.info(
-                        "%s - %s - %s", client_socket.getpeername()[0], url, first_line
-                    )
-                with open(self.html_403, "r", encoding="utf-8") as f:
-                    custom_403_page = f.read()
-                response = (
-                    f"HTTP/1.1 403 Forbidden\r\n"
-                    f"Content-Length: {len(custom_403_page)}\r\n"
-                    f"\r\n"
-                    f"{custom_403_page}"
-                )
-                client_socket.sendall(response.encode())
-                client_socket.close()
-                self.active_connections.pop(threading.get_ident(), None)
-                return
-        server_host, _ = parse_url(url)
+        if self._is_blocked(url):
+            self._send_403(client_socket, url, first_line)
+            return
+
+        parsed_url = urlparse(url)
+        server_host = parsed_url.hostname
         if not self.logger_config.no_logging_access:
             self.logger_config.access_logger.info(
                 "%s - %s - %s",
@@ -120,20 +166,16 @@ def handle_http_request(self, client_socket, request):
             )
 
         if self.config_custom_header and os.path.isfile(self.config_custom_header):
-            request_lines = request.decode(errors="ignore").split("\r\n")
-            request_line = request_lines[0]  # GET / HTTP/1.1
-
-            header_lines = [f"{key}: {value}" for key, value in headers.items()]
-            reconstructed_headers = "\r\n".join(header_lines)
-
-            if "\r\n\r\n" in request.decode(errors="ignore"):
-                body = request.decode(errors="ignore").split("\r\n\r\n", 1)[1]
-            else:
-                body = ""
-
-            modified_request = (
-                f"{request_line}\r\n{reconstructed_headers}\r\n\r\n{body}".encode()
+            request_text = request.decode(errors="ignore")
+            request_lines = request_text.split("\r\n")
+            headers = self._get_modified_headers(url, request_text)
+            request_line = request_lines[0]
+            body = (
+                request_text.split("\r\n\r\n", 1)[1]
+                if "\r\n\r\n" in request_text
+                else ""
             )
+            modified_request = self._rebuild_http_request(request_line, headers, body)
 
             self.forward_request_to_server(client_socket, modified_request, url)
 
@@ -149,10 +191,14 @@ def forward_request_to_server(self, client_socket, request, url):
             request (bytes): The raw HTTP request sent by the client.
             url (str): The target URL from the HTTP request.
         """
-        if self.proxy_enable:
-            server_host, server_port = self.proxy_host, self.proxy_port
+        if self.proxy_config.enable:
+            server_host, server_port = self.proxy_config.host, self.proxy_config.port
         else:
-            server_host, server_port = parse_url(url)
+            parsed_url = urlparse(url)
+            server_host = parsed_url.hostname
+            server_port = parsed_url.port or (
+                443 if parsed_url.scheme == "https" else 80
+            )
         thread_id = threading.get_ident()
 
         if thread_id in self.active_connections: