chisup/app/chis/auth.py

import time
from dataclasses import dataclass
from http.cookies import SimpleCookie
from typing import Mapping, Optional, Sequence
from urllib.parse import quote, urlencode, urlparse

import requests

from app.chis.crypto import sm2_encrypt


DEFAULT_ALLOWED_ROLES = ("责任医生助理", "责任医生")


class ChisLoginError(Exception):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code


@dataclass
class ChisLoginSession:
    cookie: str
    jsessionid: str
    role_id: str
    role_name: str
    user_name: str
    roles_result: dict
    apps_result: dict


class RequestsTransport:
    def __init__(self, proxy="", timeout=(10, 30)):
        self.timeout = timeout
        self.session = requests.Session()
        if proxy:
            self.session.proxies.update({"http": proxy, "https": proxy})

    def post(self, url, headers, payload):
        response = self.session.post(
            url=url,
            headers=headers,
            json=payload,
            timeout=self.timeout,
        )
        response_headers = dict(response.headers)
        if response.status_code != 200:
            return {
                "code": response.status_code,
                "msg": response.text,
                "data": None,
            }, response_headers
        if "json" in response.headers.get("Content-Type", ""):
            return response.json(), response_headers
        return response.text, response_headers


class ChisLoginClient:
    def __init__(
        self,
        base_url,
        public_key,
        transport=None,
        encrypt=None,
        now_millis=None,
        allowed_roles: Sequence[str] = DEFAULT_ALLOWED_ROLES,
    ):
        self.base_url = _normalize_chis_base_url(base_url)
        self.public_key = public_key
        self.transport = transport or RequestsTransport()
        self.encrypt = encrypt or sm2_encrypt
        self.now_millis = now_millis or _now_millis
        self.allowed_roles = tuple(allowed_roles)
        self.host = urlparse(self.base_url).netloc
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML,like Gecko) Chrome/95.0.4638.54 Safari/537.36",
            "encoding": "utf-8",
            "Content-Type": "application/json",
            "Accept": "application/json",
            "Host": self.host,
            "Referer": f"{self.base_url}/index.html",
        }

    @classmethod
    def from_config(cls, config):
        return cls(
            base_url=config.get("CHIS_BASE_URL", ""),
            public_key=config.get("CHIS_PUBLIC_KEY", ""),
            transport=RequestsTransport(proxy=config.get("CHIS_PROXY", "")),
        )

    def login(self, username, password):
        roles_result, jsessionid = self.get_my_roles(username, password)
        token = self._select_role(roles_result)
        apps_result, jsessionid = self.get_my_apps(
            username=username,
            password=password,
            role_id=str(token["id"]),
            jsessionid=jsessionid,
        )
        if not jsessionid or len(jsessionid) < 3:
            raise ChisLoginError("chis_login_failed", "CHIS login did not return JSESSIONID")

        role_name = token["roleName"]
        cookie = self._build_cookie(username, role_name, jsessionid)
        return ChisLoginSession(
            cookie=cookie,
            jsessionid=jsessionid,
            role_id=str(token["id"]),
            role_name=role_name,
            user_name=token.get("userName", ""),
            roles_result=roles_result,
            apps_result=apps_result,
        )

    def get_my_roles(self, username, password):
        payload = {
            "url": "logon/myRoles",
            "uid": username,
            "pwd": self.encrypt(self.public_key, password),
            "d": self.encrypt(self.public_key, str(self.now_millis())),
        }
        result, response_headers = self.transport.post(
            url=f"{self.base_url}/logon/myRoles",
            headers=self._headers_without_cookie(),
            payload=payload,
        )
        self._ensure_chis_ok(result, "chis_login_failed", "CHIS myRoles failed")
        jsessionid = _extract_jsessionid(response_headers)
        if not jsessionid:
            raise ChisLoginError("chis_login_failed", "CHIS myRoles did not return JSESSIONID")
        return result, jsessionid

    def get_my_apps(self, username, password, role_id, jsessionid):
        query = urlencode(
            {
                "urt": role_id,
                "uid": username,
                "pwd": self.encrypt(self.public_key, password),
                "deep": "3",
                "d": self.encrypt(self.public_key, str(self.now_millis())),
            }
        )
        headers = self._headers_without_cookie()
        headers["Cookie"] = f"JSESSIONID={jsessionid}"
        result, response_headers = self.transport.post(
            url=f"{self.base_url}/logon/myApps?{query}",
            headers=headers,
            payload={"url": f"logon/myApps?{query}", "httpMethod": "POST"},
        )
        self._ensure_chis_ok(result, "chis_login_failed", "CHIS myApps failed")
        return result, _extract_jsessionid(response_headers) or jsessionid

    def get_lander_info(self, cookie):
        headers = self._headers_without_cookie()
        headers["Cookie"] = cookie
        payload = {
            "serviceId": "chis.myPageService",
            "serviceAction": "getLanderInfo",
            "method": "execute",
        }
        result, _ = self.transport.post(
            url=f"{self.base_url}/*.jsonRequest?",
            headers=headers,
            payload=payload,
        )
        self._ensure_chis_ok(result, "chis_session_invalid", "CHIS session invalid")
        return result

    def _select_role(self, roles_result):
        tokens = roles_result.get("body", {}).get("tokens", [])
        for token in tokens:
            if token.get("roleName") in self.allowed_roles:
                return token
        roles = " / ".join(self.allowed_roles)
        raise ChisLoginError("chis_role_not_allowed", f"CHIS account has no allowed role: {roles}")

    def _headers_without_cookie(self):
        return {key: value for key, value in self.headers.items() if key.lower() != "cookie"}

    def _build_cookie(self, username, role_name, jsessionid):
        cookie_map = {
            "JSESSIONID": jsessionid,
            "sessionId1": jsessionid,
            username: f"{quote(role_name)}@photo/default.jpg",
        }
        return "; ".join([f"{key}={value}" for key, value in cookie_map.items()])

    def _ensure_chis_ok(self, result, code, prefix):
        if not isinstance(result, dict):
            raise ChisLoginError(code, f"{prefix}: invalid response")
        if result.get("code") != 200:
            message = result.get("msg") or result.get("message") or "unknown error"
            raise ChisLoginError(code, f"{prefix}: {message}")


def _normalize_chis_base_url(base_url):
    if not base_url:
        raise ChisLoginError("chis_config_missing", "CHIS_BASE_URL is not configured")
    normalized = base_url.rstrip("/")
    if "://" not in normalized:
        normalized = f"http://{normalized}"
    if normalized.endswith("/chis"):
        return normalized
    return f"{normalized}/chis"


def _now_millis():
    return int(time.time() * 1000)


def _extract_jsessionid(headers: Mapping[str, str]) -> Optional[str]:
    for key, value in headers.items():
        if key.lower() != "set-cookie":
            continue
        cookie = SimpleCookie()
        cookie.load(value)
        if "JSESSIONID" in cookie:
            return cookie["JSESSIONID"].value
        for item in value.split(";"):
            if "JSESSIONID" in item:
                return item.split("=", 1)[1].strip()
    return None