pxil-selenium-automation/tests/captcha_solver.py

"""
CAPTCHA Solver for PXIL Login using EasyOCR.
Extracts base64 from input's CSS background-image, runs EasyOCR OCR, fills the input.
"""
import re
import base64
import time
from typing import Optional

from PIL import Image
import cv2
import numpy as np
from selenium.webdriver.common.by import By
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.remote.webelement import WebElement
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import StaleElementReferenceException, TimeoutException


# Initialize EasyOCR once (module level)
# First run downloads models, then works offline
try:
    import easyocr
    _EASYOCR_READER = easyocr.Reader(['en'], gpu=False)
except ImportError:
    _EASYOCR_READER = None


# Character confusion correction map for common OCR errors
CONFUSION_MAP = {
    '0': 'O', 'O': '0', 'D': '0',
    '1': 'I', 'I': '1', 'L': '1', 'l': '1',
    '5': 'S', 'S': '5',
    '8': 'B', 'B': '8',
    '2': 'Z', 'Z': '2',
    '6': 'G', 'G': '6',
    '9': 'g', 'g': '9',
    'Q': 'O',
    'U': 'V', 'V': 'U',
    'n': 'h', 'h': 'n',
    'r': 'f', 'f': 'r',
    'c': 'e', 'e': 'c',
    'o': '0', '0': 'o',
    's': '5', '5': 's',
    'z': '2', '2': 'z',
    'i': '1', '1': 'i',
    'l': '1', '1': 'l',
}


def find_visible_captcha_input(driver: WebDriver) -> WebElement:
    """Find the visible captcha input element (name='captcha')."""
    wait = WebDriverWait(driver, 15)
    inputs = wait.until(EC.presence_of_all_elements_located((By.NAME, "captcha")))
    for inp in inputs:
        if inp.is_displayed():
            return inp
    raise RuntimeError("No visible captcha input found")


def extract_captcha_base64(captcha_input: WebElement) -> bytes:
    """Extract base64 image data from input's style attribute. Returns raw JPEG bytes."""
    style = captcha_input.get_attribute("style") or ""
    match = re.search(r'background:\s*url\(["\']?data:image/[^;]+;base64,([^"\')]+)["\']?\)', style)
    if not match:
        raise RuntimeError(f"No base64 background-image found in style: {style[:200]}")
    
    b64_data = match.group(1)
    if not b64_data:
        raise RuntimeError("Base64 data is empty (captcha not loaded yet)")
    
    b64_data += "=" * ((4 - len(b64_data) % 4) % 4)  # Fix padding
    return base64.b64decode(b64_data)


def wait_for_captcha_load(driver: WebDriver, captcha_input: WebElement, timeout: int = 15) -> bool:
    """Wait for captcha base64 to be populated with actual image data."""
    end_time = time.time() + timeout
    while time.time() < end_time:
        try:
            style = captcha_input.get_attribute("style") or ""
            match = re.search(r'background:\s*url\(["\']?data:image/[^;]+;base64,([^"\')]+)["\']?\)', style)
            if match and match.group(1):
                return True
        except StaleElementReferenceException:
            captcha_input = find_visible_captcha_input(driver)
        time.sleep(0.3)
    return False


def preprocess_captcha_image(image_bytes: bytes) -> Image.Image:
    """Preprocess image for better OCR accuracy - gentle preprocessing."""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)  # Load as color (BGR)
    
    if img is None:
        raise RuntimeError("Failed to decode image bytes")
    
    # Convert BGR to RGB
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # Convert to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    
    # Simple threshold - keep it simple
    _, gray = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY)
    
    # Light denoise
    gray = cv2.fastNlMeansDenoising(gray, h=3)
    
    # Very light morphological close - only connect broken parts
    kernel = np.ones((1, 1), np.uint8)
    gray = cv2.morphologyEx(gray, cv2.MORPH_CLOSE, kernel)
    
    # Resize 6x for better recognition (larger for 6-char captcha)
    gray = cv2.resize(gray, None, fx=6, fy=6, interpolation=cv2.INTER_CUBIC)
    
    # Convert back to RGB (3-channel) for EasyOCR
    rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
    
    return Image.fromarray(rgb)


def apply_confusion_corrections(text: str) -> str:
    """Apply character confusion corrections based on common OCR errors."""
    corrected = []
    for ch in text:
        corrected.append(CONFUSION_MAP.get(ch, ch))
    return ''.join(corrected)


def ocr_captcha(image: Image.Image) -> str:
    """Run OCR on preprocessed captcha image using EasyOCR."""
    if _EASYOCR_READER is None:
        raise RuntimeError("EasyOCR not installed. Run: pip install easyocr")
    
    img_array = np.array(image)
    
    # EasyOCR with alphanumeric allowlist for 6-char captcha
    results = _EASYOCR_READER.readtext(
        img_array,
        detail=0,
        allowlist='ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789',
        width_ths=0.5,
        paragraph=False
    )
    
    # Combine all detected text
    result = ''.join(results)
    
    # Clean result - keep only alphanumeric, max 6 chars
    cleaned = ''.join(filter(str.isalnum, result))[:6]
    
    print(f"  DEBUG OCR: result='{result}', cleaned='{cleaned}', len={len(cleaned)}")
    return cleaned


def find_refresh_button(driver: WebDriver, captcha_input: WebElement) -> Optional[WebElement]:
    """Find captcha refresh button near the captcha input."""
    parent = captcha_input.find_element(By.XPATH, "..")
    for sibling in parent.find_elements(By.XPATH, "*"):
        onclick = (sibling.get_attribute("onclick") or "").lower()
        class_attr = (sibling.get_attribute("class") or "").lower()
        id_attr = (sibling.get_attribute("id") or "").lower()
        if any(kw in onclick for kw in ["refresh", "captcha", "reload", "newcaptcha"]):
            return sibling
        if any(kw in class_attr for kw in ["refresh", "captcha-refresh", "reload"]):
            return sibling
        if any(kw in id_attr for kw in ["refresh", "captcha", "reload"]):
            return sibling
    
    # Fallback: search globally for refresh-like elements
    try:
        candidates = driver.find_elements(By.XPATH, "//*[self::a or self::button or self::img or self::span][contains(translate(@onclick, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), 'refresh') or contains(translate(@class, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), 'refresh') or contains(translate(@title, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'abcdefghijklmnopqrstuvwxyz'), 'refresh')]")
        for cand in candidates:
            if cand.is_displayed():
                return cand
    except:
        pass
    return None


def wait_for_new_captcha(driver: WebDriver, captcha_input: WebElement, timeout: int = 10) -> bool:
    """Wait for captcha background-image to change (new challenge loaded)."""
    old_style = captcha_input.get_attribute("style") or ""
    end_time = time.time() + timeout
    while time.time() < end_time:
        try:
            new_style = captcha_input.get_attribute("style") or ""
            if new_style != old_style and "data:image" in new_style:
                return True
        except StaleElementReferenceException:
            captcha_input = find_visible_captcha_input(driver)
            old_style = captcha_input.get_attribute("style") or ""
        time.sleep(0.3)
    return False


def solve_and_fill_captcha(driver: WebDriver, max_retries: int = 5, attempt_callback=None) -> bool:
    """
    Main entry point: solve captcha and fill input.
    Returns True if filled successfully, False if max retries exceeded.
    Each retry gets a fresh captcha via refresh button.
    """
    for attempt in range(1, max_retries + 1):
        captcha_text = ""
        try:
            captcha_input = find_visible_captcha_input(driver)
            
            # Wait for captcha base64 to be populated (async load)
            if not wait_for_captcha_load(driver, captcha_input):
                print(f"  [Attempt {attempt}] Captcha failed to load in time, retrying...")
                if attempt < max_retries:
                    refresh_btn = find_refresh_button(driver, captcha_input)
                    if refresh_btn:
                        refresh_btn.click()
                        wait_for_new_captcha(driver, captcha_input)
                if attempt_callback:
                    attempt_callback(attempt, "", "", "not_loaded")
                continue
            
            image_bytes = extract_captcha_base64(captcha_input)
            pil_image = preprocess_captcha_image(image_bytes)
            captcha_text = ocr_captcha(pil_image)
            
            if not captcha_text or len(captcha_text) < 6:
                print(f"  [Attempt {attempt}] OCR weak result: '{captcha_text}' (len={len(captcha_text)}), retrying...")
                if attempt < max_retries:
                    refresh_btn = find_refresh_button(driver, captcha_input)
                    if refresh_btn:
                        refresh_btn.click()
                        wait_for_new_captcha(driver, captcha_input)
                if attempt_callback:
                    attempt_callback(attempt, captcha_text, "", "ocr_failed")
                continue
            
            # Apply confusion corrections to get best variant
            corrected = apply_confusion_corrections(captcha_text)
            
            # Generate case variants for case-sensitive captcha (mixed case + numbers)
            variants = [
                captcha_text,      # as-is (from OCR)
                corrected,         # with confusion corrections
                captcha_text.upper(),     # all uppercase
                captcha_text.capitalize(), # first letter uppercase
                captcha_text.title(),     # each word capitalized
                captcha_text.swapcase(),  # swap case
            ]
            
            # Add variations with first char variations
            if len(captcha_text) >= 1:
                variants.append(captcha_text[0].upper() + captcha_text[1:])
                variants.append(captcha_text[0].lower() + captcha_text[1:])
            if len(captcha_text) >= 2:
                variants.append(captcha_text[:2].upper() + captcha_text[2:])
                variants.append(captcha_text[:2].lower() + captcha_text[2:])
            if len(captcha_text) >= 3:
                variants.append(captcha_text[:3].upper() + captcha_text[3:])
                variants.append(captcha_text[:3].lower() + captcha_text[3:])
            
            # Remove duplicates while preserving order, keep only 6-char variants
            seen = set()
            unique_variants = []
            for v in variants:
                if v not in seen and len(v) == 6:
                    seen.add(v)
                    unique_variants.append(v)
            variants = unique_variants
            
            if attempt_callback:
                screenshot_path = attempt_callback(attempt, captcha_text, "", "before_entry") or ""
            
            success = False
            # Try ALL variants before giving up on this attempt
            for variant in variants:
                if len(variant) == 6:
                    print(f"  [Attempt {attempt}] Trying variant: '{variant}'")
                    captcha_input.clear()
                    captcha_input.send_keys(variant)
                    success = True
                    if attempt_callback:
                        attempt_callback(attempt, captcha_text, variant, "filled")
                    break
            
            if not success:
                print(f"  [Attempt {attempt}] No valid 6-char variant found for '{captcha_text}', retrying...")
                if attempt < max_retries:
                    refresh_btn = find_refresh_button(driver, captcha_input)
                    if refresh_btn:
                        refresh_btn.click()
                        wait_for_new_captcha(driver, captcha_input)
                continue
            
            # Click submit
            submit_btn = WebDriverWait(driver, 10).until(
                EC.element_to_be_clickable((By.XPATH, "//button[@value='Login']"))
            )
            submit_btn.click()
            
            # Wait a bit for response
            time.sleep(2)
            
            # Check for invalid captcha error using the exact xpath
            try:
                # First handle any alerts that might be open
                try:
                    alert = driver.switch_to.alert
                    alert_text = alert.text
                    print(f"  DEBUG: Alert detected after submit: {alert_text}")
                    if "logout" in alert_text.lower() or "session" in alert_text.lower() or "previous" in alert_text.lower():
                        alert.accept()
                        print("  DEBUG: Accepted session/logout alert")
                        time.sleep(2)
                except:
                    pass
                
                invalid_captcha_elem = WebDriverWait(driver, 3).until(
                    EC.visibility_of_element_located((
                        By.XPATH,
                        "//div[@id='error']//strong[contains(normalize-space(), 'Invalid Captcha')]",
                    ))
                )
                print(f"  [Attempt {attempt}] Invalid captcha detected: {invalid_captcha_elem.text}")
                if attempt < max_retries:
                    refresh_btn = find_refresh_button(driver, captcha_input)
                    if refresh_btn:
                        refresh_btn.click()
                        wait_for_new_captcha(driver, captcha_input)
                    continue
            except TimeoutException:
                # No invalid captcha error - likely success
                print(f"  [Attempt {attempt}] No invalid captcha error detected, assuming success")
                return True
            except Exception as e:
                # Handle other exceptions like NoSuchElementException
                print(f"  [Attempt {attempt}] Error checking invalid captcha: {type(e).__name__}: {e}")
                if attempt < max_retries:
                    refresh_btn = find_refresh_button(driver, captcha_input)
                    if refresh_btn:
                        refresh_btn.click()
                        wait_for_new_captcha(driver, captcha_input)
                    continue
            
        except Exception as e:
            print(f"  [Attempt {attempt}] Error: {e}")
            if attempt < max_retries:
                refresh_btn = find_refresh_button(driver, captcha_input)
                if refresh_btn:
                    refresh_btn.click()
                    wait_for_new_captcha(driver, captcha_input)
    
    return False