BIPAuto/utils/auth.py

"""
Authentication module - Responsible for Yonyou BIP system login and logout operations
"""

import os
from pathlib import Path
from playwright.sync_api import Playwright, Browser, BrowserContext, Page, Frame

# Load environment variables
from dotenv import load_dotenv

# Get project root directory
BASE_DIR = Path(__file__).resolve().parent.parent

# Load .env file
env_path = BASE_DIR / ".env"
load_dotenv(env_path)

# Configure Playwright browser path
browsers_path = os.getenv("PLAYWRIGHT_BROWSERS_PATH")
if browsers_path:
    os.environ["PLAYWRIGHT_BROWSERS_PATH"] = browsers_path


def login(
    playwright: Playwright,
    username: str = None,
    password: str = None,
    url: str = None,
    headless: bool = None,
    ignore_https_errors: bool = None,
    verbose: bool = True,
) -> tuple[Browser, BrowserContext, Page, Frame]:
    """
    Login to Yonyou BIP system

    Args:
        playwright: Playwright instance
        username: Username (default read from environment variable ERP_USERNAME)
        password: Password (default read from environment variable ERP_PASSWORD)
        url: Login page URL (default read from environment variable ERP_URL)
        headless: Whether to use headless mode (default read from environment variable ERP_HEADLESS)
        ignore_https_errors: Whether to ignore HTTPS errors (default read from environment variable ERP_IGNORE_HTTPS_ERRORS)
        verbose: Whether to print detailed logs

    Returns:
        tuple: (browser, context, page, main_frame)
            - browser: Browser instance
            - context: Browser context
            - page: Page object
            - main_frame: Main iframe after login (forwardFrame)
    """
    import time

    # Read default configuration from environment variables
    username = username or os.getenv("ERP_USERNAME")
    password = password or os.getenv("ERP_PASSWORD")

    # URL handling: read from environment variable or parameter
    if not url:
        url = os.getenv("ERP_URL")
        if url and not url.endswith("login/main/index.html"):
            url = url.rstrip("/") + "/yonbip/resources/uap/rbac/login/main/index.html"

    if not url:
        raise ValueError("URL must be provided either as parameter or through ERP_URL environment variable")

    # headless parameter handling
    if headless is None:
        headless_str = os.getenv("ERP_HEADLESS", "false").lower()
        headless = headless_str in ("true", "1", "yes")

    # ignore_https_errors parameter handling
    if ignore_https_errors is None:
        ignore_https_errors_str = os.getenv("ERP_IGNORE_HTTPS_ERRORS", "true").lower()
        ignore_https_errors = ignore_https_errors_str in ("true", "1", "yes")

    # Launch browser
    browser = playwright.chromium.launch(headless=headless)

    # Create context
    context = browser.new_context(ignore_https_errors=ignore_https_errors)

    # Create page
    page = context.new_page()

    # Navigate to login page
    page.goto(url)

    # Extract main iframe (forwardFrame)
    main_frame = page.locator("#forwardFrame").content_frame

    # Fill username
    main_frame.get_by_role("textbox", name="用户名").fill(username)

    # Fill password
    main_frame.get_by_role("textbox", name="密码").fill(password)

    # Click login button
    main_frame.get_by_role("button", name="登录").click()

    # Handle force login popup (check if exists)
    confirm_btn = main_frame.get_by_role("button", name="确定")
    if confirm_btn.count() > 0:
        confirm_btn.click()
        if verbose:
            print("Force login detected")
    else:
        if verbose:
            print("Normal login")

    return browser, context, page, main_frame


def logout(main_frame: Frame, verbose: bool = True) -> None:
    """
    Execute account logout

    Args:
        main_frame: Main iframe object (forwardFrame)
        verbose: Whether to print detailed logs
    """
    import time

    if verbose:
        print("Clicking account menu button...")

    # Element 1: Account menu button (logo icon)
    main_frame.get_by_role("img", name="logo").click()

    # Wait for menu to appear
    time.sleep(1)

    if verbose:
        print("Clicking logout button...")

    # Element 2: "Logout" button
    main_frame.get_by_text("退出登录").click()

    # Wait for confirmation dialog to appear
    time.sleep(1)

    if verbose:
        print("Waiting for logout confirmation dialog...")

    # Element 3: Logout confirmation dialog (check if appears)
    try:
        confirm_text = main_frame.get_by_text("退出确定要退出当前账号吗？")
        if verbose:
            print("Found confirmation dialog, clicking confirm button")

        # Element 4: Confirm button
        main_frame.get_by_role("button", name="确定(Y)").click()
    except:
        if verbose:
            print("Confirmation dialog not found, may have auto-logged out")

    time.sleep(2)  # Wait for logout to complete


def close_session(browser: Browser, context: BrowserContext) -> None:
    """
    Close browser session

    Args:
        browser: Browser instance
        context: Browser context
    """
    # Check if auto-close browser is enabled
    auto_close = os.getenv("ERP_AUTO_CLOSE_BROWSER", "true").lower() in ("true", "1", "yes")

    if auto_close:
        context.close()
        browser.close()
    else:
        print("Note: Browser not auto-closed (ERP_AUTO_CLOSE_BROWSER=false)")