run_wrapper.py

#!/usr/bin/env python3
"""
VPS wrapper for the LinkedIn connection bot.

Handles:
- File-based OTP flow (for headless VPS where terminal input is unavailable)
- Chrome option patching (removes problematic --remote-debugging-pipe)
- Chrome profile lifecycle (clean create, consistent path, cleanup on exit)
- Runs run.py via subprocess
"""
import sys
import os
import time
import shutil
import signal
import subprocess
import builtins

PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
OTP_REQUEST_FILE = os.path.join(PROJECT_DIR, 'otp_request.txt')
OTP_RESPONSE_FILE = os.path.join(PROJECT_DIR, 'otp_response.txt')
CHROME_PROFILE = os.path.join(PROJECT_DIR, '.chrome-profile')

# -------------------------------------------------------------------------
# Chrome process management
# -------------------------------------------------------------------------

def kill_chrome_processes():
    """Kill any existing Chrome/Chromium processes."""
    try:
        subprocess.run(['pkill', '-f', 'chrome'], capture_output=True)
        subprocess.run(['pkill', '-f', 'chromium'], capture_output=True)
        time.sleep(1)
    except Exception:
        pass

def cleanup_chrome_profile():
    """Remove the Chrome profile directory."""
    if os.path.exists(CHROME_PROFILE):
        try:
            shutil.rmtree(CHROME_PROFILE)
            print(f"[WRAPPER] Cleaned up Chrome profile: {CHROME_PROFILE}", flush=True)
        except Exception as e:
            print(f"[WRAPPER] Warning: could not remove Chrome profile: {e}", flush=True)

# -------------------------------------------------------------------------
# OTP flow (file-based for VPS)
# -------------------------------------------------------------------------

_driver_ref = [None]

def request_otp(prompt="Enter verification code"):
    """Write OTP request file and wait for response file (up to 600s)."""
    print(f"[OTP] Requesting input: {prompt}", flush=True)
    with open(OTP_REQUEST_FILE, 'w') as f:
        f.write(f"{time.time()}\n{prompt}")
    for _ in range(600):
        if os.path.exists(OTP_RESPONSE_FILE):
            with open(OTP_RESPONSE_FILE, 'r') as f:
                response = f.read().strip()
            os.remove(OTP_RESPONSE_FILE)
            try:
                os.remove(OTP_REQUEST_FILE)
            except OSError:
                pass
            print("[OTP] Got response", flush=True)
            return response
        time.sleep(1)
    try:
        os.remove(OTP_REQUEST_FILE)
    except OSError:
        pass
    return ""

def _submit_otp_code(driver, code):
    """Submit an OTP code into the verification form."""
    from selenium.webdriver.common.by import By
    try:
        input_field = driver.find_element(By.CSS_SELECTOR,
            'input[name="pin"], input#input__email_verification_pin, input[type="text"]')
        input_field.clear()
        input_field.send_keys(code)
        time.sleep(1)
        submit_btn = driver.find_element(By.CSS_SELECTOR,
            'button[type="submit"], button#email-pin-submit-button')
        submit_btn.click()
        time.sleep(5)
        print(f"[VERIFY] Code submitted. URL: {driver.current_url}", flush=True)
    except Exception as e:
        print(f"[VERIFY] Error submitting code: {e}", flush=True)


def smart_input(prompt=""):
    """
    Replacement for builtins.input() on VPS.
    When a verification challenge is detected:
      1. Try automatic OTP retrieval from email (IMAP)
      2. Fall back to file-based OTP exchange if IMAP is not configured
    """
    driver = _driver_ref[0]
    if driver is not None:
        try:
            current_url = driver.current_url
            if 'checkpoint' in current_url or 'challenge' in current_url:
                print("[VERIFY] Verification challenge detected!", flush=True)
                screenshot_path = os.path.join(PROJECT_DIR, 'verify-screenshot.png')
                driver.save_screenshot(screenshot_path)

                # Strategy 1: Auto-read OTP from email via IMAP
                code = None
                try:
                    from otp_email import fetch_otp
                    code = fetch_otp()
                except ImportError:
                    print("[VERIFY] otp_email module not available, skipping auto-read.", flush=True)
                except Exception as e:
                    print(f"[VERIFY] Auto-read failed: {e}", flush=True)

                # Strategy 2: File-based OTP exchange (fallback)
                if not code:
                    print("[VERIFY] Falling back to file-based OTP exchange...", flush=True)
                    code = request_otp("LinkedIn verification code needed. Check your email.")

                if code:
                    _submit_otp_code(driver, code)
                return ""
        except Exception as e:
            print(f"[VERIFY] Check failed: {e}", flush=True)

    # Non-challenge input() calls also try file-based exchange
    return request_otp(prompt)

# -------------------------------------------------------------------------
# Monkey-patches (applied before importing run.py)
# -------------------------------------------------------------------------

# Replace input() with file-based OTP handler
builtins.input = smart_input

# Patch Chrome options to remove problematic flags
import selenium.webdriver.chrome.options
_OrigOpts = selenium.webdriver.chrome.options.Options
class PatchedOptions(_OrigOpts):
    def add_argument(self, arg):
        if '--remote-debugging-pipe' in arg:
            return
        super().add_argument(arg)
selenium.webdriver.chrome.options.Options = PatchedOptions

# Patch Chrome init to force our profile path and capture driver ref
import selenium.webdriver
_orig_init = selenium.webdriver.Chrome.__init__
def _patched_init(self, *args, **kwargs):
    if 'options' in kwargs and kwargs['options'] is not None:
        opts = kwargs['options']
        if not any('--user-data-dir' in a for a in opts.arguments):
            opts.add_argument(f'--user-data-dir={CHROME_PROFILE}')
    _orig_init(self, *args, **kwargs)
    _driver_ref[0] = self
selenium.webdriver.Chrome.__init__ = _patched_init

# -------------------------------------------------------------------------
# Main execution
# -------------------------------------------------------------------------

if __name__ == '__main__':
    os.chdir(PROJECT_DIR)

    # 1. Kill existing Chrome processes
    print("[WRAPPER] Killing existing Chrome processes...", flush=True)
    kill_chrome_processes()

    # 2. Clean existing Chrome profile
    cleanup_chrome_profile()
    os.makedirs(CHROME_PROFILE, exist_ok=True)
    print(f"[WRAPPER] Created fresh Chrome profile: {CHROME_PROFILE}", flush=True)

    # 3. Bypass env_bootstrap in subprocess (venv already active on VPS)
    sys.modules['env_bootstrap'] = type(sys)('env_bootstrap')
    sys.modules['env_bootstrap'].ensure_venv = lambda: None

    # 4. Run the bot
    exit_code = 1
    try:
        # Import and execute run.py directly (patches must be in same process)
        with open('run.py', 'r') as f:
            source = f.read()
        # Remove the env_bootstrap import lines since we bypassed it
        source = source.replace('from env_bootstrap import ensure_venv\nensure_venv()', '')
        exec(compile(source, 'run.py', 'exec'), {'__name__': '__main__', '__file__': 'run.py'})
        exit_code = 0
    except Exception as e:
        print(f"[WRAPPER] Bot execution error: {e}", flush=True)
        raise
    finally:
        # 5. Cleanup: kill Chrome, delete profile, clean verification emails
        print("[WRAPPER] Cleaning up...", flush=True)
        kill_chrome_processes()
        cleanup_chrome_profile()
        # Remove OTP files
        for f in [OTP_REQUEST_FILE, OTP_RESPONSE_FILE]:
            try:
                os.remove(f)
            except OSError:
                pass
        # Clean up LinkedIn verification emails from inbox
        try:
            from otp_email import cleanup_linkedin_emails
            cleanup_linkedin_emails()
        except Exception:
            pass
        print(f"[WRAPPER] Done. Exit code: {exit_code}", flush=True)