Merge pull request #509 from AutomationSolutionz/mitim-proxy

Proxy implementation for capturing browser generated traffic

Author: Sakib75

Framework/Built_In_Automation/Sequential_Actions/action_declarations/common.py

@@ -131,6 +131,8 @@
 
     {"name": "connect to google service client",            "function": "connect_to_google_service_account",     "screenshot": "none" },
     {"name": "upload to google storage bucket",             "function": "upload_to_google_storage_bucket",       "screenshot": "none" },
+
+    {"name": "proxy server",                                "function": "proxy_server",                         "screenshot": "none"}
 
 ) # yapf: disable
 

Framework/Built_In_Automation/Sequential_Actions/common_functions.py

@@ -6740,3 +6740,72 @@ def stop_ssh_tunnel(data_set):
 
     return "passed"
 
+@logger
+def proxy_server(data_set):
+    import os
+    sModuleInfo = inspect.currentframe().f_code.co_name + " : " + MODULE_NAME
+
+    proxy_var = None
+    action = None
+    port = 8080
+    for left, mid, right in data_set:
+        if left.lower().strip() == 'action':
+            action = 'start' if right.lower().strip() == 'start' else 'stop'
+        if left.lower().strip() == 'port':
+            port = int(right.strip())
+        if left.lower().strip() == 'proxy server':
+            proxy_var = right.strip()
+    
+    if action == None:
+        CommonUtil.ExecLog(sModuleInfo, "Incorrect dataset", 3)
+        return "zeuz_failed"
+
+
+    if action == 'start':        
+        CommonUtil.ExecLog(sModuleInfo, f"{action.capitalize()}ing proxy server on port {port}", 1)
+
+        proxy_log_dir = Path(sr.Get_Shared_Variables("zeuz_download_folder")).parent / 'proxy_log'
+        os.makedirs(proxy_log_dir, exist_ok=True)
+        mitm_proxy_path = Path(__file__).parent / "mitm_proxy.py"
+        output_file_path = proxy_log_dir / 'mitm.log'  # Output file to save the logs
+        CommonUtil.ExecLog(sModuleInfo, f"Proxy Log file: {output_file_path}", 1)
+
+        captured_network_file_path = proxy_log_dir / 'captured_network_data.csv'
+        CommonUtil.ExecLog(sModuleInfo, f"Captured Network file: {output_file_path}", 1)
+        # Open the output file in append mode
+        with open(r'{}'.format(output_file_path), 'a') as output_file:
+            # Start the subprocess
+            process = subprocess.Popen(
+                [
+                    "mitmdump",
+                    "-s",
+                    f"{mitm_proxy_path}",
+                    "-p",
+                    str(port),
+                    "--set",
+                    f"output_file_path={captured_network_file_path}",
+                ],
+                stdout=output_file,  # Redirect stdout to the file
+                stderr=output_file,  # Redirect stderr to the file
+            )
+            
+        pid = process.pid
+        CommonUtil.mitm_proxy_pids.append(pid)
+        CommonUtil.ExecLog(sModuleInfo, f"Started process with PID: {pid}", 1)
+
+        sr.Set_Shared_Variables(proxy_var, {"pid":pid,"captured_network_file_path":captured_network_file_path,"log_file":output_file_path})
+        return "passed"
+    else:
+        import signal
+        
+        if CommonUtil.mitm_proxy_pids:
+            try:
+                pid = CommonUtil.mitm_proxy_pids[0]
+                os.kill(pid, signal.SIGTERM)
+                CommonUtil.ExecLog(sModuleInfo,f"Process with PID {pid} has been terminated.",1)
+                CommonUtil.mitm_proxy_pids.pop()
+            except OSError as e:
+                CommonUtil.ExecLog(sModuleInfo,f"Error: {e}", 3)
+
+        CommonUtil.ExecLog(sModuleInfo, f"{action.capitalize()}ing proxy server on port {port}", 1)
+        return "passed"

Framework/Built_In_Automation/Sequential_Actions/mitm_proxy.py

@@ -0,0 +1,81 @@
+from mitmproxy import http, ctx
+import time
+import csv
+import os
+
+# Initialize a list to track request details
+requests_data = []
+
+
+def load(l):
+    # Define the custom option `output_file_path`
+    ctx.options.add_option("output_file_path", str, "", "Path to output CSV file")
+
+
+def request(flow: http.HTTPFlow) -> None:
+    # Capture request data when it's made
+    start_time = time.time()
+    requests_data.append({
+        'url': flow.request.url,
+        'start_time': start_time,
+        'status_code': None,
+        'end_time': None,
+        'duration': None,
+        'content_length': None
+    })
+
+def response(flow: http.HTTPFlow) -> None:
+    output_file_path = ctx.options.output_file_path
+    create_file_if_not_exists(output_file_path)
+
+    # print("Flow", flow)
+    # print("Response", flow.response)
+
+    res = flow.response
+    end_time = time.time()
+
+    # Find the matching request based on the URL
+    for req in requests_data:
+        if req['url'] == flow.request.url:
+            req['status_code'] = res.status_code
+            req['end_time'] = end_time
+            req['duration'] = end_time - req['start_time']
+            req['content_length'] = len(res.content)
+            break
+
+    # Create a list to hold the captured details
+    captured_details = [
+        flow.request.url,
+        res.status_code,
+        req.get('duration', None),
+        len(res.content),
+        end_time
+    ]
+
+    # Append the captured details as a row in the CSV file
+    with open(output_file_path, 'a', newline='') as csvfile:
+        writer = csv.writer(csvfile)
+        writer.writerow(captured_details)  # Write CSV row
+
+    # Optionally print captured details for console output
+    print(f"Captured: {captured_details}")
+
+def create_file_if_not_exists(filepath):
+    """
+    Check if the output CSV file exists.
+    If it does not exist, create the file and add csv headers.
+    """
+
+    if not os.path.exists(filepath):
+        with open(filepath, "w", newline="") as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(
+                [
+                    "url",
+                    "status_code",
+                    "duration_in_seconds",
+                    "content_length_in_bytes",
+                    "timestamp",
+                ]
+            )
+        print(f"Created output file: {filepath}")

Framework/Utilities/CommonUtil.py

@@ -189,6 +189,8 @@
 global_sleep = {"selenium":{}, "appium":{}, "windows":{}, "desktop":{}}
 zeuz_disable_var_print = {}
 
+mitm_proxy_pids = []
+
 def clear_performance_metrics():
     """reset everything to initial value"""
     global browser_perf, action_perf, step_perf, test_case_perf, perf_test_perf, api_performance_data, load_testing, processed_performance_data

Framework/test.py

@@ -0,0 +1,57 @@
+import os
+import subprocess
+from Utilities import CommonUtil
+from pathlib import Path
+
+OUTPUT_FILEPATH = "/Users/sakib/Documents/zeuz/Zeuz_Python_Node/AutomationLog/debug_sakib_dd446e56-a_AaRlG/session_1/TEST-10894/zeuz_download_folder/proxy_log/mitm.log"
+CAPTURED_CSV_FILEPATH = "/Users/sakib/Documents/zeuz/Zeuz_Python_Node/AutomationLog/debug_sakib_dd446e56-a_AaRlG/session_1/TEST-10894/zeuz_download_folder/proxy_log/captured_network_data.csv"
+MITM_PROXY_PATH = "/Users/sakib/Documents/zeuz/Zeuz_Python_Node/Framework/Built_In_Automation/Sequential_Actions/mitm_proxy.py"
+PORT = 8080 
+
+
+print(f"Starting proxy server on port {PORT}")
+
+print(f"MITM Proxy path: {MITM_PROXY_PATH}")
+print(f"Proxy Log file: {OUTPUT_FILEPATH}")
+
+print(f"Captured Network file: {CAPTURED_CSV_FILEPATH}")
+
+# Open the output file in append mode
+with open(OUTPUT_FILEPATH, 'a') as output_file:
+    # Start the subprocess
+    process = subprocess.Popen(
+        ['mitmdump', '-s', MITM_PROXY_PATH, '-w', str(CAPTURED_CSV_FILEPATH), '-p', str(PORT)],
+        stdout=output_file,  # Redirect stdout to the file
+        stderr=output_file   # Redirect stderr to the file
+    )
+    
+pid = process.pid
+
+# Assuming CommonUtil.mitm_proxy_pids is a list, make sure it's initialized properly
+if not hasattr(CommonUtil, 'mitm_proxy_pids'):
+    CommonUtil.mitm_proxy_pids = []
+
+CommonUtil.mitm_proxy_pids.append(pid)
+
+import time
+time.sleep(2)
+
+# Verify if the service is running on the specified port
+def verify_port_in_use(port):
+    if os.name == 'posix':  # macOS/Linux
+        result = subprocess.run(['lsof', '-i', f':{port}'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        return result.stdout
+    elif os.name == 'nt':  # Windows
+        result = subprocess.run(['netstat', '-aon'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        return result.stdout if f':{port}' in result.stdout else ''
+
+# Check if the port is in use
+port_status = verify_port_in_use(PORT)
+
+if port_status:
+    print(f"Service is running on port {PORT}:\n{port_status}")
+else:
+    print(f"Service is NOT running on port {PORT}. Check if the subprocess started correctly.")
+
+# Prevent the script from exiting immediately
+input("Press Enter to exit...\n")

requirements-linux.txt

@@ -56,6 +56,6 @@ jinja2
 pandas
 pyperclip
 thefuzz
-backports-datetime-fromisoformat; python_version < '3.11'
 genson
-google-cloud-storage
+google-cloud-storage
+mitmproxy

requirements-mac.txt

@@ -59,7 +59,7 @@ configobj
 jinja2
 pandas
 pyperclip
-backports-datetime-fromisoformat; python_version < '3.11'
 thefuzz
 genson
-google-cloud-storage
+google-cloud-storage
+mitmproxy

requirements-win.txt

@@ -69,7 +69,7 @@ configobj
 jinja2
 pandas
 pyperclip
-backports-datetime-fromisoformat; python_version < '3.11'
 thefuzz
 genson
-google-cloud-storage
+google-cloud-storage
+mitmproxy

requirements.txt

@@ -51,6 +51,6 @@ configobj
 boto3
 pandas
 pyperclip
-backports-datetime-fromisoformat; python_version < '3.11'
 genson
-google-cloud-storage
+google-cloud-storage
+mitmproxy