Merge pull request #49 from VectorlyApp/fix_unexpected_kill

Fix unexpected Chrome session termination

Author: rayruizhiliao

quickstart.py

@@ -11,8 +11,6 @@
 import subprocess
 import shutil
 import json
-import atexit
-import signal
 from pathlib import Path
 from typing import Optional
 import requests
@@ -29,9 +27,6 @@
 CDP_CAPTURES_DIR = Path("./cdp_captures")
 DISCOVERY_OUTPUT_DIR = Path("./routine_discovery_output")
 
-# Global variable to track Chrome process for cleanup
-_chrome_process: Optional[subprocess.Popen] = None
-
 
 def print_colored(text: str, color: str = NC) -> None:
     """Print colored text."""
@@ -292,55 +287,6 @@ def launch_chrome(port: int) -> Optional[subprocess.Popen]:
         return None
 
 
-def cleanup_chrome(process: Optional[subprocess.Popen], port: int) -> None:
-    """Clean up Chrome process if it was launched by this script."""
-    if process is None:
-        return
-    
-    try:
-        # Check if Chrome is still running on the port
-        if not check_chrome_running(port):
-            return
-        
-        print()
-        print_colored("🧹 Cleaning up Chrome...", YELLOW)
-        
-        # Try graceful termination first
-        try:
-            if platform.system() == "Windows":
-                # On Windows with CREATE_NEW_PROCESS_GROUP, we need to kill the process group
-                process.terminate()
-                time.sleep(1)
-                if process.poll() is None:
-                    process.kill()
-            else:
-                process.terminate()
-                time.sleep(1)
-                if process.poll() is None:
-                    process.kill()
-            
-            # Wait a bit for Chrome to close
-            process.wait(timeout=3)
-            print_colored("✅ Chrome closed successfully", GREEN)
-        except subprocess.TimeoutExpired:
-            # Force kill if it didn't terminate
-            try:
-                process.kill()
-                process.wait(timeout=2)
-                print_colored("✅ Chrome force-closed", GREEN)
-            except Exception:
-                pass
-        except Exception as e:
-            # Process might already be dead
-            if process.poll() is not None:
-                print_colored("✅ Chrome already closed", GREEN)
-            else:
-                print_colored(f"⚠️  Error closing Chrome: {e}", YELLOW)
-    except Exception:
-        # Silently fail during cleanup
-        pass
-
-
 def run_command(cmd: list[str], description: str) -> bool:
     """Run a command and return True if successful."""
     try:
@@ -360,8 +306,6 @@ def run_command(cmd: list[str], description: str) -> bool:
 
 def main():
     """Main workflow."""
-    global _chrome_process
-    
     # Use local variables that can be updated
     cdp_captures_dir = CDP_CAPTURES_DIR
     discovery_output_dir = DISCOVERY_OUTPUT_DIR
@@ -370,7 +314,7 @@ def main():
     print_colored("║         Web Hacker - Quickstart Workflow                   ║", BLUE)
     print_colored("╚════════════════════════════════════════════════════════════╝", BLUE)
     print()
-    
+
     # Pipeline overview
     print_colored("Web-hacker Pipeline:", BLUE)
     print()
@@ -390,19 +334,6 @@ def main():
         print_colored(f"✅ Chrome is already running in debug mode on port {PORT}", GREEN)
     else:
         chrome_process = launch_chrome(PORT)
-        # Store globally for cleanup
-        _chrome_process = chrome_process
-        # Register cleanup function if we launched Chrome
-        if chrome_process is not None:
-            atexit.register(cleanup_chrome, chrome_process, PORT)
-            # Also register signal handlers for graceful shutdown
-            def signal_handler(signum, frame):
-                cleanup_chrome(chrome_process, PORT)
-                sys.exit(0)
-            signal.signal(signal.SIGINT, signal_handler)
-            # SIGTERM may not be available on all platforms
-            if hasattr(signal, 'SIGTERM'):
-                signal.signal(signal.SIGTERM, signal_handler)
 
     print()
 
@@ -456,14 +387,6 @@ def signal_handler(signum, frame):
         run_command(monitor_cmd, "monitoring")
         print()
 
-    # Close Chrome before Step 3 if we launched it
-    if chrome_process is not None:
-        cleanup_chrome(chrome_process, PORT)
-        atexit.unregister(cleanup_chrome)
-        chrome_process = None
-        _chrome_process = None
-        print()
-    
     # Step 3: Discover
     print_colored("Step 3: Discovering routine from captured data...", GREEN)
 
@@ -475,9 +398,11 @@ def signal_handler(signum, frame):
         return
 
     skip = input("   Skip discovery step? (y/n): ").strip().lower()
+    
     if skip == 'y':
         # Use default directory when skipping - user can specify routine path in step 4 if needed
         discovery_output_dir = DISCOVERY_OUTPUT_DIR
+        
         print_colored("⏭️  Skipping discovery step.", GREEN)
         print_colored(f"   Using default discovery output directory: {discovery_output_dir.resolve()}", GREEN)
         print()
@@ -489,8 +414,6 @@ def signal_handler(signum, frame):
         if new_output_dir:
             discovery_output_dir = Path(new_output_dir)
             print_colored(f"✅ Using discovery output directory: {discovery_output_dir}", GREEN)
-        else:
-            discovery_output_dir = DISCOVERY_OUTPUT_DIR
 
         # Check if routine already exists
         routine_file = discovery_output_dir / "routine.json"
@@ -621,8 +544,4 @@ def signal_handler(signum, frame):
     except KeyboardInterrupt:
         print()
         print_colored("⚠️  Interrupted by user.", YELLOW)
-        # Clean up Chrome if we launched it
-        if _chrome_process is not None:
-            cleanup_chrome(_chrome_process, PORT)
-        sys.exit(0)
-
+        sys.exit(0)

web_hacker/cdp/cdp_session.py

@@ -42,10 +42,6 @@ def __init__(
         self.clear_cookies = clear_cookies
         self.clear_storage = clear_storage
 
-        # Connection state tracking
-        self._connection_lost = False
-        self._connection_lost_lock = threading.Lock()
-        
         # Response tracking for synchronous commands
         self.pending_responses = {}
         self.response_lock = threading.Lock()
@@ -71,15 +67,8 @@ def __init__(
 
     def send(self, method, params=None):
         """Send CDP command and return sequence ID."""
-        if self._connection_lost:
-            raise ConnectionError("WebSocket connection is closed")
         self.seq += 1
-        try:
-            self.ws.send(json.dumps({"id": self.seq, "method": method, "params": params or {}}))
-        except (websocket.WebSocketConnectionClosedException, OSError, ConnectionError) as e:
-            with self._connection_lost_lock:
-                self._connection_lost = True
-            raise ConnectionError(f"WebSocket connection lost: {e}")
+        self.ws.send(json.dumps({"id": self.seq, "method": method, "params": params or {}}))
         return self.seq
 
     def send_and_wait(self, method, params=None, timeout=10):
@@ -111,14 +100,6 @@ def send_and_wait(self, method, params=None, timeout=10):
 
     def setup_cdp(self, navigate_to=None):
         """Setup CDP domains and configuration."""
-        # Enable Target domain to receive target lifecycle events (if on browser WebSocket)
-        # Note: This may not work on tab WebSockets, but that's okay - we'll catch disconnections
-        try:
-            self.send("Target.setDiscoverTargets", {"discover": True})
-        except Exception as e:
-            # This is expected to fail on tab WebSockets - Target domain is browser-level
-            logger.debug(f"Could not enable Target domain (expected on tab WebSockets): {e}")
-        
         # Enable basic domains
         self.send("Page.enable")
         self.send("Runtime.enable")
@@ -167,14 +148,6 @@ def setup_cdp(self, navigate_to=None):
 
     def handle_message(self, msg):
         """Handle incoming CDP message by delegating to appropriate monitors."""
-        # Check for target lifecycle events (tab closure)
-        method = msg.get("method")
-        if method == "Target.targetDestroyed":
-            logger.info("Tab was closed. Connection will be lost. Saving assets...")
-            with self._connection_lost_lock:
-                self._connection_lost = True
-            return
-        
         # Try network monitor first
         if self.network_monitor.handle_network_message(msg, self):
             return
@@ -229,84 +202,44 @@ def _handle_command_reply(self, msg):
 
         return False
 
-    def _generate_assets(self):
-        """Generate all monitoring assets. Works even if connection is lost."""
-        try:
-            # Final cookie sync using native CDP (only if connection is still alive)
-            if not self._connection_lost:
-                try:
-                    self.storage_monitor.monitor_cookie_changes(self)
-                except Exception as e:
-                    logger.debug(f"Could not sync cookies (connection may be lost): {e}")
-        except Exception as e:
-            logger.debug(f"Error in cookie sync: {e}")
+    def run(self):
+        """Main message processing loop."""
+        logger.info("Blocking trackers & capturing network/storage… Press Ctrl+C to stop.")
 
-        # Consolidate all transactions into a single JSON file (works with cached data)
         try:
+            while True:
+                msg = json.loads(self.ws.recv())
+                self.handle_message(msg)
+        except KeyboardInterrupt:
+            logger.info("\nStopped. Saving assets...")
+            # Final cookie sync using native CDP (no delay needed)
+            self.storage_monitor.monitor_cookie_changes(self)
+            
+            # Consolidate all transactions into a single JSON file
             consolidated_path = f"{self.output_dir}/consolidated_transactions.json"
             self.network_monitor.consolidate_transactions(consolidated_path)
-        except Exception as e:
-            logger.warning(f"Could not consolidate transactions: {e}")
-        
-        # Generate HAR file from consolidated transactions (works with cached data)
-        try:
+            
+            # Generate HAR file from consolidated transactions
             har_path = f"{self.output_dir}/network.har"
             self.network_monitor.generate_har_from_transactions(har_path, "Web Hacker Session")
-        except Exception as e:
-            logger.warning(f"Could not generate HAR file: {e}")
-        
-        # Consolidate all interactions into a single JSON file (works with cached data)
-        try:
+            
+            # Consolidate all interactions into a single JSON file
             interaction_dir = self.paths.get('interaction_dir', f"{self.output_dir}/interaction")
             consolidated_interactions_path = os.path.join(interaction_dir, "consolidated_interactions.json")
             self.interaction_monitor.consolidate_interactions(consolidated_interactions_path)
-        except Exception as e:
-            logger.warning(f"Could not consolidate interactions: {e}")
-    
-    def run(self):
-        """Main message processing loop."""
-        logger.info("Blocking trackers & capturing network/storage… Press Ctrl+C to stop.")
-        
-        try:
-            while True:
-                try:
-                    msg = json.loads(self.ws.recv())
-                    self.handle_message(msg)
-                    
-                    # Check if connection was lost due to tab closure
-                    if self._connection_lost:
-                        logger.info("Tab closed. Saving assets...")
-                        break
-                except (websocket.WebSocketConnectionClosedException, OSError, ConnectionError) as e:
-                    # WebSocket connection lost (tab closed, browser closed, etc.)
-                    logger.info(f"Connection lost: {e}. Saving assets...")
-                    with self._connection_lost_lock:
-                        self._connection_lost = True
-                    break
-                except json.JSONDecodeError as e:
-                    logger.warning(f"Failed to parse message: {e}")
-                    continue
-        except KeyboardInterrupt:
-            logger.info("\nStopped by user. Saving assets...")
         finally:
-            # Always generate assets, even if connection is lost
-            self._generate_assets()
-            
-            # Close WebSocket if still open
             try:
-                if self.ws and not self._connection_lost:
-                    self.ws.close()
-            except Exception:
+                self.ws.close()
+            except:
                 pass
 
     def get_monitoring_summary(self):
         """Get summary of all monitoring activities."""
-        # Trigger final cookie check using native CDP (only if connection is still alive)
-        if not self._connection_lost:
-            try:
-                self.storage_monitor.monitor_cookie_changes(self)
-            except Exception as e:
-                logger.debug(f"Could not sync cookies for summary: {e}")
+        # Trigger final cookie check using native CDP (no delay needed)
+        try:
+            self.storage_monitor.monitor_cookie_changes(self)
+        except:
+            pass
 
         storage_summary = self.storage_monitor.get_storage_summary()
         network_summary = self.network_monitor.get_network_summary()
@@ -316,4 +249,4 @@ def get_monitoring_summary(self):
             "network": network_summary,
             "storage": storage_summary,
             "interaction": interaction_summary,
-        }
+        }