add feature custom header

Author: 6C656C65

.gitignore

@@ -171,8 +171,8 @@ cython_debug/
 .pypirc
 
 # Config file
-config/*.txt
-!config/*.example.txt
+config/*
+!config/*.example*
 
 # Certs folder
 certs/*.key

README.md

@@ -36,6 +36,7 @@
 - Support distant (http) blacklist
 - Shortcuts
 - Cancel inspection on bank site
+- Custom header
 
 ## 📦 **Installation**
 
@@ -86,7 +87,6 @@ If you encounter any problems, or if you want to use the program in a particular
 - Benchmark
 - Admin mode, statistiques and real time request
 - Image slim wihtout admin interface
-- Custom header
 - Fix HSTS
 - Wiki ssl inspection (volume certs/ca...)
 

config.ini.example

@@ -20,6 +20,7 @@ blocked_url = config/blocked_url.txt
 
 [Options]
 shortcuts = config/shortcuts.txt
+custom_header = config/custom_header.json
 
 [Security]
 ssl_inspect = false

config/custom_header.example.json

@@ -0,0 +1,9 @@
+{
+    "https://example.com": {
+      "X-Custom-Header": "HeaderValue",
+      "X-Another-Header": "AnotherValue"
+    },
+    "https://example2.com": {
+      "X-Different-Header": "DifferentValue"
+    }
+  }

pyproxy.py

@@ -70,6 +70,11 @@
         type=str,
         help="Path to the text file containing the list of shortcuts"
     )
+    parser.add_argument(
+        "--custom-header",
+        type=str,
+        help="Path to the json file containing the list of custom headers"
+    )
     parser.add_argument("--no-logging-access", action="store_true", help="Disable access logging")
     parser.add_argument("--no-logging-block", action="store_true", help="Disable block logging")
     parser.add_argument("--ssl-inspect", action="store_true", help="Enable SSL inspection")
@@ -133,6 +138,11 @@
         if args.blocked_url
         else config.get('Options', 'shortcuts', fallback="config/shortcuts.txt")
     )
+    custom_header = (
+        args.blocked_url
+        if args.blocked_url
+        else config.get('Options', 'custom_header', fallback="config/custom_header.json")
+    )
     no_logging_access = (
         args.no_logging_access
         if args.no_logging_access
@@ -184,6 +194,7 @@
         blocked_sites=blocked_sites,
         blocked_url=blocked_url,
         shortcuts=shortcuts,
+        custom_header=custom_header,
         inspect_ca_cert=inspect_ca_cert,
         inspect_ca_key=inspect_ca_key,
         inspect_certs_folder=inspect_certs_folder,

utils/cancel_inspect.py

@@ -59,9 +59,7 @@ def file_monitor() -> None:
         try:
             while True:
                 new_cancel_inspect = load_cancel_inspect(cancel_inspect_path)
-                cancel_inspect_data.clear()
-                cancel_inspect_data.extend(new_cancel_inspect)
-
+                cancel_inspect_data[:] = new_cancel_inspect
                 time.sleep(5)
         except (IOError, ValueError) as e:
             print(f"File monitor error: {e}")

utils/custom_header.py

@@ -0,0 +1,79 @@
+"""
+custom_header.py
+
+This module contains functions and a process to load and monitor custom header entries.
+It reads a file containing custom header data and checks whether specific entries exist in that file.
+The file is monitored in a background thread for live updates.
+
+Functions:
+- load_custom_header: Loads custom header entries from a file into a list.
+- custom_header_process: Process that listens for header-like entries and checks
+  if they exist in the custom header list.
+"""
+
+import multiprocessing
+import time
+import sys
+import threading
+import json
+
+def load_custom_header(custom_header_path: str) -> dict:
+    """
+    Loads custom header entries from a file into a list.
+    
+    Args:
+        custom_header_path (str): The path to the file containing the custom headers.
+    
+    Returns:
+        dict: A dictionary containing the custom header data loaded from the file.
+    """
+    with open(custom_header_path, 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+# pylint: disable=too-many-locals
+def custom_header_process(
+    queue: multiprocessing.Queue,
+    result_queue: multiprocessing.Queue,
+    custom_header_path: str
+) -> None:
+    """
+    Process that monitors the custom header file and checks if received entries exist in it.
+    
+    Args:
+        queue (multiprocessing.Queue): A queue to receive header-like entries to check.
+        result_queue (multiprocessing.Queue): A queue to send back True/False depending on match.
+        custom_header_path (str): Path to the file containing custom header entries.
+    """
+    manager = multiprocessing.Manager()
+    custom_header_data = manager.dict(
+        load_custom_header(custom_header_path)
+    )
+
+    error_event = threading.Event()
+
+    def file_monitor() -> None:
+        try:
+            while True:
+                new_custom_header = load_custom_header(custom_header_path)
+                custom_header_data.clear()
+                custom_header_data.update(new_custom_header)
+                time.sleep(5)
+        except (IOError, ValueError) as e:
+            print(f"File monitor error: {e}")
+            error_event.set()
+
+    monitor_thread = threading.Thread(target=file_monitor, daemon=True)
+    monitor_thread.start()
+
+    while True:
+        if error_event.is_set():
+            print("Error detected in file monitor thread, terminating process.")
+            sys.exit(1)
+
+        try:
+            url = queue.get()
+            headers = custom_header_data.get(url, {})
+            result_queue.put(headers)
+
+        except KeyboardInterrupt:
+            break

utils/proxy.py

@@ -23,6 +23,7 @@
 from utils.filter import filter_process
 from utils.shortcuts import shortcuts_process
 from utils.cancel_inspect import cancel_inspect_process
+from utils.custom_header import custom_header_process
 from utils.logger import configure_file_logger, configure_console_logger
 
 class ProxyServer:
@@ -35,7 +36,7 @@ class ProxyServer:
     # pylint: disable=too-many-locals
     def __init__(self, host, port, debug, access_log, block_log,
                  html_403, no_filter, filter_mode, no_logging_access, no_logging_block, ssl_inspect,
-                 blocked_sites, blocked_url, shortcuts, inspect_ca_cert,
+                 blocked_sites, blocked_url, shortcuts, custom_header, inspect_ca_cert,
                  inspect_ca_key, inspect_certs_folder, cancel_inspect):
         """
         Initializes the ProxyServer instance with the provided configurations.
@@ -57,10 +58,14 @@ def __init__(self, host, port, debug, access_log, block_log,
         self.cancel_inspect_proc = None
         self.cancel_inspect_queue = multiprocessing.Queue()
         self.cancel_inspect_result_queue = multiprocessing.Queue()
+        self.custom_header_proc = None
+        self.custom_header_queue = multiprocessing.Queue()
+        self.custom_header_result_queue = multiprocessing.Queue()
         self.console_logger = configure_console_logger()
         self.config_blocked_sites = blocked_sites
         self.config_blocked_url = blocked_url
         self.config_shortcuts = shortcuts
+        self.config_custom_header = custom_header
         self.config_cancel_inspect = cancel_inspect
         self.config_inspect_cert = inspect_ca_cert
         self.config_inspect_key = inspect_ca_key
@@ -91,6 +96,7 @@ def start(self):
             self.console_logger.debug("[*] blocked_sites = %s", self.config_blocked_sites)
             self.console_logger.debug("[*] blocked_url = %s", self.config_blocked_url)
             self.console_logger.debug("[*] shortcuts = %s", self.config_shortcuts)
+            self.console_logger.debug("[*] custom_header = %s", self.config_custom_header)
             self.console_logger.debug("[*] inspect_ca_cert = %s", self.config_inspect_cert)
             self.console_logger.debug("[*] inspect_ca_key = %s", self.config_inspect_key)
             self.console_logger.debug(
@@ -162,6 +168,18 @@ def start(self):
             self.cancel_inspect_proc.start()
             self.console_logger.debug("[*] Starting the cancel inspection process...")
 
+        if self.config_custom_header and os.path.isfile(self.config_custom_header):
+            self.custom_header_proc = multiprocessing.Process(
+                target=custom_header_process,
+                args=(
+                    self.custom_header_queue,
+                    self.custom_header_result_queue,
+                    self.config_custom_header
+                )
+            )
+            self.custom_header_proc.start()
+            self.console_logger.debug("[*] Starting the custom header process...")
+
         server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         server.bind(self.host_port)
         server.listen(10)
@@ -213,6 +231,13 @@ 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 = self.extract_headers(request.decode(errors='ignore'))
+            self.custom_header_queue.put(url)
+            new_headers = self.custom_header_result_queue.get()
+            headers.update(new_headers)
+            print("headers", headers)
+
         if self.config_shortcuts:
             domain, _ = self.parse_url(url)
             self.shortcuts_queue.put(domain)
@@ -259,7 +284,26 @@ def handle_http_request(self, client_socket, request):
                 f"http://{server_host}",
                 first_line
             )
-        self.forward_request_to_server(client_socket, request, url)
+
+        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()
+
+            # 5. Envoyer au serveur
+            self.forward_request_to_server(client_socket, modified_request, url)
+
+        else:
+            self.forward_request_to_server(client_socket, request, url)
 
     def forward_request_to_server(self, client_socket, request, url):
         """
@@ -320,6 +364,24 @@ def parse_url(self, url):
 
         return server_host, server_port
 
+    def extract_headers(self, request_str):
+        """
+        Extracts the HTTP headers from a raw HTTP request string.
+
+        Args:
+            request_str (str): The full HTTP request as a decoded string.
+
+        Returns:
+            dict: A dictionary containing the HTTP header fields as key-value pairs.
+        """
+        headers = {}
+        lines = request_str.split("\n")[1:]
+        for line in lines:
+            if line.strip():
+                key, value = line.split(":", 1)
+                headers[key.strip()] = value.strip()
+        return headers
+
     # pylint: disable=too-many-locals,too-many-statements,too-many-branches,too-many-nested-blocks
     def handle_https_connection(self, client_socket, first_line):
         """