Peng_ACL_AutoProtect/auth_session.py

"""Combined authentication and in-memory session state.

This module merges `session_state.py` and `login_post.py` into a single
convenient file. It exposes the session-state helpers:

- set_sess_key(key: str) -> None
- get_sess_key() -> Optional[str]
- clear_sess_key() -> None

and the `login` function which performs the POST request and stores the
`sess_key` (if found) into the in-memory session state.

It also provides a `main()` function intended for CLI use.
"""
from __future__ import annotations

import json
import re
import sys
from pathlib import Path
from threading import Lock
from typing import Optional, Tuple

import requests

from log_util import get_logger, log_request


ROOT = Path(__file__).resolve().parent
CONFIG_PATH = ROOT / "config.json"


# --- simple thread-safe in-memory session state ---
_lock = Lock()
_sess_key: Optional[str] = None


def set_sess_key(key: str) -> None:
    """Set global sess_key (thread-safe)."""
    global _sess_key
    with _lock:
        _sess_key = key


def get_sess_key() -> Optional[str]:
    """Get current sess_key, or None if not set (thread-safe)."""
    with _lock:
        return _sess_key


def clear_sess_key() -> None:
    """Clear current sess_key (thread-safe)."""
    global _sess_key
    with _lock:
        _sess_key = None


# --- login/request helpers ---
def load_config() -> dict:
    if not CONFIG_PATH.exists():
        raise FileNotFoundError(f"配置文件未找到: {CONFIG_PATH}")
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        return json.load(f)


DEFAULT_PAYLOAD = {
    "username": "xiaobai",
    "passwd": "dc81b4427df07fd6b3ebcb05a7b34daf",
    "pass": "c2FsdF8xMXhpYW9iYWku",
    "remember_password": "",
}


def _extract_sess_cookie_from_response(resp: requests.Response) -> Optional[str]:
    """Try to extract sess_key cookie string from response.

    Returns a cookie string like 'sess_key=...;' if found, else None.
    """
    try:
        sess_val = resp.cookies.get("sess_key")
        if sess_val:
            return f"sess_key={sess_val};"
        set_cookie = resp.headers.get("Set-Cookie", "")
        m = re.search(r"(sess_key=[^;]+;?)", set_cookie)
        if m:
            return m.group(1)
    except Exception:
        return None
    return None


def login(payload: dict | None = None, timeout: int = 10) -> Tuple[requests.Response, Optional[str]]:
    """Send login POST. Returns (response, sess_cookie_or_none).

    The function will also set the in-memory sess_key if it can be extracted.
    """
    cfg = load_config()
    base = cfg.get("base_url", "").rstrip("/")
    url = f"{base}/Action/login"
    data = payload or DEFAULT_PAYLOAD
    logger = get_logger("login_post")
    logger.debug(f"准备发送请求，URL: {url}")
    resp = requests.post(url, json=data, timeout=timeout)

    sess_cookie = _extract_sess_cookie_from_response(resp)
    if sess_cookie:
        # Write sess_key into in-memory session state (strip trailing semicolon)
        cookie_val = sess_cookie.split("=", 1)[1].rstrip(";")
        try:
            set_sess_key(cookie_val)
            logger.info("sess_key 已保存到内存状态")
        except Exception:
            logger.exception("保存 sess_key 失败")

    # Log request/response but do not fail on logging errors
    try:
        log_request(logger, "login", url, data, resp)
    except Exception:
        logger.exception("记录请求/响应失败")

    return resp, sess_cookie


def main() -> None:
    logger = get_logger("main")
    try:
        resp, sess_cookie = login()
    except FileNotFoundError as e:
        logger.error(f"配置错误: {e}")
        sys.exit(2)
    except requests.RequestException as e:
        logger.error(f"请求失败: {e}")
        sys.exit(1)

    logger.info(f"状态: {resp.status_code}")
    content_type = resp.headers.get("Content-Type", "")
    if "application/json" in content_type:
        try:
            pretty = json.dumps(resp.json(), ensure_ascii=False, indent=2)
            logger.info(f"响应 JSON:\n{pretty}")
            print(pretty)
        except ValueError:
            logger.info(f"响应文本: {resp.text}")
            print(resp.text)
    else:
        logger.info(f"响应文本: {resp.text}")
        print(resp.text)

    if sess_cookie:
        logger.info(f"提取到 sess_key: {sess_cookie}")
        print(f"sess_key: {sess_cookie}")


if __name__ == "__main__":
    main()