scheduler/backend/scanner.py

import os
import re
import threading

try:
    import smbclient  # type: ignore[import-not-found]
except ModuleNotFoundError:
    smbclient = None

from db import DEVICE_DUT, DEVICE_REF, mark_tests_completed, mark_overdue_as_rerun

_SMB_SESSIONS = set()
_SCAN_STATE_LOCK = threading.Lock()
_ACTIVE_SCAN_COUNT = 0
_RESULT_TEST_ID_PATTERN = re.compile(r"(?:COE|P2P|P3P)(?:RX|TX)?[A-Z]{2}\d{3}", re.IGNORECASE)


def _normalize_smb_credentials(smb_credentials=None):
    username = ""
    password = ""
    domain = ""

    if smb_credentials:
        username = str(smb_credentials.get("username", "")).strip()
        password = str(smb_credentials.get("password", ""))
        domain = str(smb_credentials.get("domain", "")).strip()

    if not username:
        username = os.getenv("SMB_USERNAME", "").strip()
    if not password:
        password = os.getenv("SMB_PASSWORD", "")
    if not domain:
        domain = os.getenv("SMB_DOMAIN", "").strip()

    if username and domain and "\\" not in username and "@" not in username:
        username = f"{domain}\\{username}"

    return username, password


def _extract_test_id_from_result_dir_name(dir_name):
    if not dir_name:
        return None

    match = _RESULT_TEST_ID_PATTERN.search(str(dir_name).upper())
    if match:
        return match.group(0)
    return None


def _build_match_tokens(parsed):
    tokens = set()

    def _add(value):
        if value in (None, ""):
            return
        tokens.add(str(value).upper())

    _add(parsed.get("interference"))
    _add(parsed.get("device"))
    _add(parsed.get("rotation"))
    _add(parsed.get("test_point"))
    _add(parsed.get("rssi"))
    _add(parsed.get("station"))
    _add(parsed.get("band"))
    _add(parsed.get("channel"))
    _add(parsed.get("bandwidth"))
    _add(parsed.get("direction"))
    _add(parsed.get("throttled"))
    _add(parsed.get("test_id"))

    for bw in parsed.get("extra_bandwidths") or []:
        _add(bw)

    # Add numeric alias for devices, e.g. CGW453 -> 453.
    device = (parsed.get("device") or "").upper()
    device_digits = re.sub(r"\D", "", device)
    if device_digits:
        tokens.add(device_digits)

    # SP can appear as flag or specific token variant (e.g., SP40).
    if parsed.get("sp"):
        tokens.add("SP")
        tokens.add(str(parsed.get("sp")).upper())

    return tokens


def _rule_matches_target(rule, parsed_tokens):
    # Rule parts are AND-ed: CGW453_P3P_ROT2 => device AND interference AND rotation.
    # Support _, -, or spaces as condition separators.
    parts = [p for p in re.split(r"[_\-\s]+", rule) if p]
    if not parts:
        return False

    def _part_matches(part):
        if part in parsed_tokens:
            return True

        # Support tag variants (e.g., BW80 should match BW80M/BW80+80 in source names).
        # Keep this conservative for very short parts to avoid overmatching.
        if len(part) >= 3:
            for token in parsed_tokens:
                if token.startswith(part) or part in token:
                    return True

        return False

    return all(_part_matches(part) for part in parts)


def _should_exclude_target(parsed, exclusions):
    parsed_tokens = _build_match_tokens(parsed)
    # Any matching rule excludes the testcase.
    return any(_rule_matches_target(rule, parsed_tokens) for rule in exclusions)


def _scan_started():
    global _ACTIVE_SCAN_COUNT
    with _SCAN_STATE_LOCK:
        _ACTIVE_SCAN_COUNT += 1


def _scan_finished():
    global _ACTIVE_SCAN_COUNT
    with _SCAN_STATE_LOCK:
        _ACTIVE_SCAN_COUNT = max(0, _ACTIVE_SCAN_COUNT - 1)


def is_scan_in_progress():
    with _SCAN_STATE_LOCK:
        return _ACTIVE_SCAN_COUNT > 0


def resolve_runtime_path(path_value):
    if not path_value:
        return path_value

    raw_path = str(path_value).strip()
    if not raw_path:
        return raw_path

    # If the path is already valid in the current runtime, keep it.
    if os.path.exists(raw_path):
        return raw_path

    # UNC/network paths are handled separately via smbclient.
    if raw_path.startswith("\\\\") or raw_path.startswith("//"):
        return raw_path

    # Map host paths (Windows or Linux) to the container mount point when running in a container.
    if os.name != "nt":
        mount_root = os.getenv("HOST_MOUNT_ROOT", "/host").strip() or "/host"
        host_root = os.getenv("HOST_BROWSE_ROOT", "").strip()

        raw_norm = raw_path.replace("\\", "/")
        if host_root:
            host_norm = host_root.replace("\\", "/").rstrip("/")
            if raw_norm.lower() == host_norm.lower() or raw_norm.lower().startswith(host_norm.lower() + "/"):
                relative = raw_norm[len(host_norm):].lstrip("/")
                if relative:
                    return os.path.join(mount_root, *relative.split("/"))
                return mount_root

    return raw_path


def _normalize_input_path(path_value):
    if not path_value:
        return path_value

    path = resolve_runtime_path(path_value)
    path = str(path).strip()

    # Normalize forward slashes to backslashes first
    path = path.replace("/", "\\")

    # Accept \\server\share style UNC paths and ensure proper escaping
    if path.startswith("\\\\"):
        # Clean up any doubled backslashes from replacement
        while "\\\\\\" in path:
            path = path.replace("\\\\\\", "\\\\")
        return path

    # Accept //server/share style and normalize to UNC for smbclient.
    if path.startswith("//"):
        path = path.lstrip("/").replace("/", "\\")
        return "\\\\" + path

    # Accept /<ipv4>/<share>/... and normalize to UNC for Linux-hosted inputs.
    if re.match(r"^\\?\d{1,3}(?:\\\.\d{1,3}){3}\\[^\\]+", path):
        if path.startswith("\\"):
            path = path.lstrip("\\")
        return "\\\\" + path

    return path


def _is_unc_path(path):
    return isinstance(path, str) and path.startswith("\\\\")


def _extract_unc_server(path):
    if not _is_unc_path(path):
        return None
    rest = path[2:]
    return rest.split("\\", 1)[0] if rest else None


def _register_smb_session_if_needed(path, smb_credentials=None):
    if not _is_unc_path(path):
        return

    if smbclient is None:
        raise ModuleNotFoundError("smbclient is required to scan UNC result paths")

    server = _extract_unc_server(path)
    if not server or server in _SMB_SESSIONS:
        return

    username, password = _normalize_smb_credentials(smb_credentials)

    if username:
        smbclient.register_session(server, username=username, password=password)
    else:
        smbclient.register_session(server)

    _SMB_SESSIONS.add(server)


def _iter_dir_entries(path, smb_credentials=None):
    path = _normalize_input_path(path)
    if _is_unc_path(path):
        if smbclient is None:
            raise ModuleNotFoundError("smbclient is required to scan UNC result paths")
        _register_smb_session_if_needed(path, smb_credentials=smb_credentials)
        return list(smbclient.scandir(path))
    return list(os.scandir(path))


def path_exists_with_smb(path, smb_credentials=None):
    """Check if a path exists, with SMB authentication for UNC paths."""
    if not path:
        return False
    
    path = _normalize_input_path(path)
    
    # For UNC paths, use SMB to check
    if _is_unc_path(path):
        try:
            _register_smb_session_if_needed(path, smb_credentials=smb_credentials)
            # Try to list entries; if successful, path exists
            _iter_dir_entries(path, smb_credentials=smb_credentials)
            return True
        except Exception:
            return False
    
    # For local paths, use standard os.path.exists
    return os.path.exists(path)


def scan_results(results_dir_dut, results_dir_ref, smb_credentials=None):
    results_dir_dut = _normalize_input_path(results_dir_dut)
    results_dir_ref = _normalize_input_path(results_dir_ref)

    if not results_dir_dut or not results_dir_ref:
        return
    
    try:
        dut_entries = [
            entry.name
            for entry in _iter_dir_entries(results_dir_dut, smb_credentials=smb_credentials)
            if entry.is_dir()
        ]
    except OSError as exc:
        print(f"[scanner] Cannot read results dir: {exc}")
        return

    print(f"[scanner] results: {len(dut_entries)} result dir(s) found in DUT results dir")

    try:
        ref_entries = [
            entry.name
            for entry in _iter_dir_entries(results_dir_ref, smb_credentials=smb_credentials)
            if entry.is_dir()
            and not entry.name.startswith("obsolete")
        ]
    except OSError as exc:
        print(f"[scanner] Cannot read results dir: {exc}")
        return
    
    print(f"[scanner] results: {len(ref_entries)} result dir(s) found in reference results dir")

    completed_batch = []
    unmatched_entries = []

    for entry_name in dut_entries:
        test_id = _extract_test_id_from_result_dir_name(entry_name)
        if test_id:
            completed_batch.append((test_id, DEVICE_DUT))
        else:
            unmatched_entries.append(entry_name)

    for entry_name in ref_entries:
        test_id = _extract_test_id_from_result_dir_name(entry_name)
        if test_id:
            completed_batch.append((test_id, DEVICE_REF))
        else:
            unmatched_entries.append(entry_name)

    if unmatched_entries:
        print(f"[scanner] skipped {len(unmatched_entries)} result dir(s) with no recognizable test id")

    if completed_batch:
        print(f"[scanner] marking {len(completed_batch)} test(s) as completed:")
        for test_id, device in completed_batch:
            print(f"  - {test_id} on {device}")
    
    updated_count = mark_tests_completed(completed_batch)
    print(f"[scanner] {updated_count} test(s) actually updated in database")
    
    newly_rerun = mark_overdue_as_rerun()
    if newly_rerun:
        print(f"[scanner] {newly_rerun} test(s) marked as rerun-required (scheduled but not completed)")