Peng_ACL_AutoProtect/login_post.py

import json
import sys
from pathlib import Path

import requests

from log_util import get_logger, log_request


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


def load_config():
    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 login(payload: dict | None = None, timeout: int = 10):
    """Send login POST. Returns requests.Response.

    payload: JSON payload to send. If None, uses DEFAULT_PAYLOAD.
    timeout: request timeout in seconds.
    """
    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)

    # Try to extract sess_key from cookies or Set-Cookie header.
    sess_cookie = None
    try:
        # requests exposes cookies in resp.cookies
        sess_val = resp.cookies.get("sess_key")
        if sess_val:
            sess_cookie = f"sess_key={sess_val};"
        else:
            set_cookie = resp.headers.get("Set-Cookie", "")
            import re

            m = re.search(r"(sess_key=[^;]+;?)", set_cookie)
            if m:
                sess_cookie = m.group(1)
    except Exception:
        # Non-fatal: if extraction fails, return None for sess_cookie
        sess_cookie = None

    # 记录请求/响应
    try:
        log_request(logger, "login", url, data, resp)
    except Exception:
        logger.exception("记录请求/响应失败")

    return resp, sess_cookie


def main():
    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}")
    # Try to pretty-print JSON if possible
    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:
        # Print the raw sess_key cookie string
        logger.info(f"提取到 sess_key: {sess_cookie}")
        print(f"sess_key: {sess_cookie}")


if __name__ == "__main__":
    main()