2fa_bot.py

import pyotp
import base64
import logging
import time
import subprocess
import os
import sys
from dataclasses import dataclass
from typing import Optional

# ================= Configuration =================
SECRET_KEY = os.environ.get("TOTP_SECRET")
X11_DISPLAY = os.environ.get("DISPLAY_NUM", ":1")
AUTOFILL_ENABLED = os.environ.get("IBKR_2FA_AUTOFILL", "yes").strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}

# Timing constants (seconds)
CHECK_INTERVAL = 3
FILL_COOLDOWN = 60
TYPE_DELAY_MS = 100
PRE_ENTER_DELAY = 1
XDOTOOL_TIMEOUT = 10
MIN_TOTP_SECONDS_REMAINING = 15
MAX_AUTOFILL_SUBMISSIONS_RAW = os.environ.get("IBKR_2FA_MAX_SUBMISSIONS", "3")
MAX_AUTOFILL_SUBMISSIONS_PER_WINDOW_RAW = os.environ.get(
    "IBKR_2FA_MAX_SUBMISSIONS_PER_WINDOW",
    "1",
)
SUBMISSION_RESET_SECONDS_RAW = os.environ.get("IBKR_2FA_SUBMISSION_RESET_SECONDS", "0")
DISMISS_LOGIN_MESSAGES = os.environ.get("IBKR_DISMISS_LOGIN_MESSAGES", "yes").strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}
DISMISS_SMALL_GATEWAY_DIALOGS = os.environ.get(
    "IBKR_DISMISS_SMALL_GATEWAY_DIALOGS",
    "yes",
).strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}

# Window titles to search for 2FA prompts. Live IBKR accounts can show mobile
# push / IB Key wording instead of the shorter TOTP-oriented prompts.
SEARCH_PATTERNS = [
    "Second Factor",
    "Challenge",
    "Security Code",
    "Enter Code",
    "IBKR Mobile",
    "IB Key",
    "Two-Factor",
    "Two Factor",
    "Verification",
    "Verify",
    "Authentication",
]
AUTH_TITLE_KEYWORDS = (
    "second factor",
    "challenge",
    "security code",
    "enter code",
    "ibkr mobile",
    "ib key",
    "two-factor",
    "two factor",
    "verification",
    "verify",
    "authentication",
)
IGNORED_TITLE_KEYWORDS = (
    "authenticating",
)
DISMISSIBLE_DIALOG_SEARCH_PATTERNS = [
    "Login Messages",
    "IBKR Gateway",
    "Gateway",
]
DISMISSIBLE_DIALOG_TITLE_KEYWORDS = (
    "login messages",
)
SMALL_GATEWAY_DIALOG_TITLE = "ibkr gateway"
SMALL_CONNECTION_DIALOG_TITLE = "gateway"
SMALL_GATEWAY_DIALOG_MAX_WIDTH = 650
SMALL_GATEWAY_DIALOG_MAX_HEIGHT = 220
# Current IBKR Gateway TOTP prompts place the code field in the upper half of
# the compact dialog. Keep the click centered on the text field instead of the
# button/link area below it.
INPUT_CLICK_POSITION = (0.50, 0.40)
# =================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("2fa_bot")

def parse_non_negative_int_env(name, raw_value):
    try:
        parsed = int(raw_value)
    except ValueError:
        log.error("%s must be an integer", name)
        sys.exit(1)
    if parsed < 0:
        log.error("%s must be non-negative", name)
        sys.exit(1)
    return parsed


MAX_AUTOFILL_SUBMISSIONS = parse_non_negative_int_env(
    "IBKR_2FA_MAX_SUBMISSIONS",
    MAX_AUTOFILL_SUBMISSIONS_RAW,
)
MAX_AUTOFILL_SUBMISSIONS_PER_WINDOW = parse_non_negative_int_env(
    "IBKR_2FA_MAX_SUBMISSIONS_PER_WINDOW",
    MAX_AUTOFILL_SUBMISSIONS_PER_WINDOW_RAW,
)
SUBMISSION_RESET_SECONDS = parse_non_negative_int_env(
    "IBKR_2FA_SUBMISSION_RESET_SECONDS",
    SUBMISSION_RESET_SECONDS_RAW,
)
autofill_submission_count = 0
autofill_limit_warned = False
window_submission_counts = {}
window_limit_warned = set()
last_submission_reset_at = time.monotonic()
dismissed_dialog_windows = set()


@dataclass(frozen=True)
class WindowCandidate:
    window_id: str
    title: str
    width: Optional[int] = None
    height: Optional[int] = None


def validate_config():
    """Validate required config at startup; exit immediately if invalid."""
    if not AUTOFILL_ENABLED:
        log.info("TOTP auto-fill is disabled; bot will only log auth popup detection")
        return
    if MAX_AUTOFILL_SUBMISSIONS < 1:
        log.error("IBKR_2FA_MAX_SUBMISSIONS must be at least 1")
        sys.exit(1)
    if MAX_AUTOFILL_SUBMISSIONS_PER_WINDOW < 1:
        log.error("IBKR_2FA_MAX_SUBMISSIONS_PER_WINDOW must be at least 1")
        sys.exit(1)
    if not SECRET_KEY:
        log.error("TOTP_SECRET not found in environment variables")
        sys.exit(1)
    try:
        base64.b32decode(SECRET_KEY.upper().replace(" ", ""), casefold=True)
    except Exception:
        log.error("TOTP_SECRET is not valid base32")
        sys.exit(1)


def get_totp():
    """Calculate 6-digit dynamic verification code using pyotp."""
    return pyotp.TOTP(SECRET_KEY).now()


def totp_seconds_remaining():
    return 30 - (int(time.time()) % 30)


def run_xdotool(args, sensitive=False):
    """Execute xdotool command on the X11 display with timeout protection."""
    env = os.environ.copy()
    env["DISPLAY"] = X11_DISPLAY
    command = ["xdotool", *args]
    try:
        return subprocess.run(
            command, env=env,
            capture_output=True, text=True,
            timeout=XDOTOOL_TIMEOUT,
        )
    except subprocess.TimeoutExpired:
        display_command = "xdotool <redacted>" if sensitive else " ".join(command)
        log.warning("xdotool command timed out: %s", display_command)
        return subprocess.CompletedProcess(command, 1, stdout="", stderr="timeout")


def get_window_title(window_id):
    res = run_xdotool(["getwindowname", window_id])
    return res.stdout.strip()


def get_window_geometry(window_id):
    res = run_xdotool(["getwindowgeometry", "--shell", window_id])
    if res.returncode != 0:
        return None, None

    values = {}
    for line in res.stdout.splitlines():
        if "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key] = value

    try:
        return int(values.get("WIDTH", "")), int(values.get("HEIGHT", ""))
    except ValueError:
        return None, None


def is_auth_candidate(title):
    normalized_title = title.lower()
    if any(keyword in normalized_title for keyword in IGNORED_TITLE_KEYWORDS):
        return False
    return any(keyword in normalized_title for keyword in AUTH_TITLE_KEYWORDS)


def is_small_gateway_dialog(title, width, height):
    if not DISMISS_SMALL_GATEWAY_DIALOGS:
        return False
    if title.lower() != SMALL_GATEWAY_DIALOG_TITLE:
        return False
    if not width or not height:
        return False
    return width <= SMALL_GATEWAY_DIALOG_MAX_WIDTH and height <= SMALL_GATEWAY_DIALOG_MAX_HEIGHT


def is_small_connection_dialog(title, width, height):
    if not DISMISS_SMALL_GATEWAY_DIALOGS:
        return False
    if title.lower() != SMALL_CONNECTION_DIALOG_TITLE:
        return False
    if not width or not height:
        return False
    return width <= SMALL_GATEWAY_DIALOG_MAX_WIDTH and height <= SMALL_GATEWAY_DIALOG_MAX_HEIGHT


def is_dismissible_dialog_candidate(title, width=None, height=None):
    normalized_title = title.lower()
    if any(keyword in normalized_title for keyword in DISMISSIBLE_DIALOG_TITLE_KEYWORDS):
        return True
    return is_small_gateway_dialog(title, width, height) or is_small_connection_dialog(title, width, height)


def find_windows_by_patterns(patterns):
    window_ids = []
    for pattern in patterns:
        res = run_xdotool(["search", "--name", pattern])
        if res.returncode != 0:
            continue
        for window_id in res.stdout.splitlines():
            if window_id and window_id not in window_ids:
                window_ids.append(window_id)
    return window_ids


def find_auth_windows():
    """Find visible IBKR authentication windows without relying on focus state."""
    window_ids = find_windows_by_patterns(SEARCH_PATTERNS)

    candidates = []
    for window_id in reversed(window_ids):
        title = get_window_title(window_id)
        if not is_auth_candidate(title):
            continue
        width, height = get_window_geometry(window_id)
        candidates.append(WindowCandidate(window_id, title, width, height))
    return candidates


def find_dismissible_dialogs():
    if not DISMISS_LOGIN_MESSAGES:
        return []

    candidates = []
    for window_id in reversed(find_windows_by_patterns(DISMISSIBLE_DIALOG_SEARCH_PATTERNS)):
        title = get_window_title(window_id)
        width, height = get_window_geometry(window_id)
        if not is_dismissible_dialog_candidate(title, width, height):
            continue
        candidates.append(WindowCandidate(window_id, title, width, height))
    return candidates


def dismiss_dialog(candidate):
    if candidate.window_id not in dismissed_dialog_windows:
        log.info(
            "Dismissing post-login dialog (id=%s, title=%r, size=%sx%s)",
            candidate.window_id,
            candidate.title,
            candidate.width or "?",
            candidate.height or "?",
        )
        dismissed_dialog_windows.add(candidate.window_id)

    run_xdotool(["windowactivate", "--sync", candidate.window_id])
    run_xdotool(["windowfocus", "--sync", candidate.window_id])
    time.sleep(0.2)
    run_xdotool(["key", "Return"])


def dismiss_post_login_dialogs():
    candidates = find_dismissible_dialogs()
    for candidate in candidates:
        dismiss_dialog(candidate)
        return True
    return False


def wait_for_fresh_totp_window():
    seconds_remaining = totp_seconds_remaining()
    if seconds_remaining > MIN_TOTP_SECONDS_REMAINING:
        return seconds_remaining

    wait_seconds = seconds_remaining + 1
    log.info(
        "TOTP period has %ss remaining; waiting %ss before submitting a fresh code",
        seconds_remaining,
        wait_seconds,
    )
    time.sleep(wait_seconds)
    return totp_seconds_remaining()


def focus_input_area(candidate):
    if not candidate.width or not candidate.height:
        return

    x = max(1, int(candidate.width * INPUT_CLICK_POSITION[0]))
    y = max(1, int(candidate.height * INPUT_CLICK_POSITION[1]))
    run_xdotool(["windowactivate", "--sync", candidate.window_id])
    run_xdotool(["windowfocus", "--sync", candidate.window_id])
    run_xdotool(["mousemove", "--window", candidate.window_id, str(x), str(y)])
    run_xdotool(["click", "1"])
    time.sleep(0.2)


def type_totp_into_active_window(code):
    """Type into the focused control; Java dialogs can ignore direct window events."""
    run_xdotool(["key", "ctrl+a", "BackSpace"])
    run_xdotool(["type", "--delay", str(TYPE_DELAY_MS), code], sensitive=True)
    time.sleep(PRE_ENTER_DELAY)
    run_xdotool(["key", "Return"])


def maybe_reset_submission_counters():
    global autofill_limit_warned
    global autofill_submission_count
    global last_submission_reset_at
    global window_limit_warned
    global window_submission_counts

    if SUBMISSION_RESET_SECONDS <= 0:
        return

    now = time.monotonic()
    if now - last_submission_reset_at < SUBMISSION_RESET_SECONDS:
        return

    if autofill_submission_count or window_submission_counts:
        log.info(
            "Resetting 2FA auto-fill counters after %ss",
            SUBMISSION_RESET_SECONDS,
        )
    autofill_submission_count = 0
    autofill_limit_warned = False
    window_submission_counts = {}
    window_limit_warned = set()
    last_submission_reset_at = now


def submit_totp(candidate):
    """Submit a TOTP code to the selected authentication popup."""
    global autofill_limit_warned
    global autofill_submission_count

    if not AUTOFILL_ENABLED:
        log.info(
            "Authentication window found (id=%s, title=%r, size=%sx%s); auto-fill disabled",
            candidate.window_id,
            candidate.title,
            candidate.width or "?",
            candidate.height or "?",
        )
        return

    maybe_reset_submission_counters()

    if autofill_submission_count >= MAX_AUTOFILL_SUBMISSIONS:
        if not autofill_limit_warned:
            log.warning(
                "Authentication window found but auto-fill submission limit reached "
                "(limit=%s); leaving window for manual handling",
                MAX_AUTOFILL_SUBMISSIONS,
            )
            autofill_limit_warned = True
        return

    window_submission_count = window_submission_counts.get(candidate.window_id, 0)
    if window_submission_count >= MAX_AUTOFILL_SUBMISSIONS_PER_WINDOW:
        if candidate.window_id not in window_limit_warned:
            log.warning(
                "Authentication window %s already received %s auto-fill submission(s); "
                "leaving it for manual handling",
                candidate.window_id,
                window_submission_count,
            )
            window_limit_warned.add(candidate.window_id)
        return

    seconds_remaining = wait_for_fresh_totp_window()
    log.info(
        "Authentication window found (id=%s, title=%r, size=%sx%s); submitting code with %ss remaining",
        candidate.window_id,
        candidate.title,
        candidate.width or "?",
        candidate.height or "?",
        seconds_remaining,
    )

    focus_input_area(candidate)
    code = get_totp()
    type_totp_into_active_window(code)

    autofill_submission_count += 1
    window_submission_counts[candidate.window_id] = window_submission_count + 1
    log.info("Auto-fill submitted, waiting for gateway response...")


def find_and_fill():
    """Search for the IBKR Gateway authentication window and auto-fill the code."""
    candidates = find_auth_windows()
    for candidate in candidates:
        submit_totp(candidate)
        return True
    return False


def main():
    validate_config()
    log.info("IBKR 2FA Bot started, monitoring display %s", X11_DISPLAY)

    while True:
        try:
            if dismiss_post_login_dialogs():
                time.sleep(1)
                continue
            if find_and_fill():
                time.sleep(FILL_COOLDOWN)
        except Exception as e:
            log.exception("Runtime exception: %s", e)

        time.sleep(CHECK_INTERVAL)


if __name__ == "__main__":
    main()