setup.py
from __future__ import annotations
import asyncio
import getpass
import os
import re
import sys
import threading
from urllib.parse import urlsplit, urlunsplit
import aiohttp
from config import DEFAULT_TELEMETRY_URL, prepare_runtime, resolve_db_path
prepare_runtime()
from db import store
from mail.imap_reader import validate_imap_cfg
from mail.oauth import (
OAuthError,
authorize_url,
device_code_wait,
exchange_code,
guess_provider,
parse_auth_code,
pkce_pair,
provider,
start_device,
start_loopback_receiver,
wait_loopback_code,
)
DISCORD_API = "https://discord.com/api/v10"
_EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")
INVITE_PERMISSIONS = (1 << 10) | (1 << 11) | (1 << 14) | (1 << 28)
SECURE_ROLE_NAME = "Secure Access"
_RESET = "\033[0m"
_BOLD = "\033[1m"
_DIM = "\033[2m"
_CYAN = "\033[36m"
_GREEN = "\033[32m"
_RED = "\033[31m"
_YELLOW = "\033[33m"
def _use_color() -> bool:
if os.environ.get("NO_COLOR"):
return False
if os.environ.get("FORCE_COLOR"):
return True
return sys.stdout.isatty()
def _paint(code: str, text: str) -> str:
if not _use_color():
return text
return f"{code}{text}{_RESET}"
def _banner(subtitle: str) -> None:
line = "─" * 46
print()
print(_paint(_CYAN, f" {line}"))
print(_paint(_BOLD, " Open AutoSecure"))
print(_paint(_DIM, f" {subtitle}"))
print(_paint(_CYAN, f" {line}"))
print()
def _step(n: int, total: int, title: str) -> None:
print()
print(_paint(_BOLD, f" Step {n}/{total}") + _paint(_DIM, f" {title}"))
def _ok(msg: str) -> None:
print(_paint(_GREEN, " ✓ ") + msg)
def _err(msg: str) -> None:
print(_paint(_RED, " ✗ ") + msg)
def _info(msg: str) -> None:
print(" " + _paint(_DIM, msg))
def _wait(msg: str) -> None:
print(_paint(_YELLOW, " … ") + msg)
def _pause(msg: str = "Press Enter to retry") -> None:
input(f" {msg} ")
def _prompt(label: str, *, default: str | None = None) -> str:
suffix = f" [{default}]" if default else ""
value = input(f" {label}{suffix}: ").strip()
return value or (default or "")
def _yes(label: str, *, default: bool = False) -> bool:
hint = "Y/n" if default else "y/N"
value = _prompt(f"{label} [{hint}]")
if not value:
return default
return value.lower() in ("y", "yes")
def _mask_proxy(url: str | None) -> str:
if not url:
return "(direct)"
parts = urlsplit(url)
if parts.username is None or parts.password is None:
return url
host = parts.hostname or ""
if ":" in host and not host.startswith("["):
host = f"[{host}]"
port = f":{parts.port}" if parts.port else ""
netloc = f"{parts.username}:***@{host}{port}"
return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
def valid_proxy_url(value: str) -> bool:
parts = urlsplit((value or "").strip())
return parts.scheme in {"http", "https"} and bool(parts.hostname)
def _print_proxy_help() -> None:
_info("Type: HTTP or HTTPS proxy (CONNECT). SOCKS4/SOCKS5 are not supported.")
_info("Format: http://[user:password@]host:port")
_info("Examples: http://10.0.0.1:8080")
_info(" http://user:pass@10.0.0.1:8080")
_info(" https://proxy.example.com:443")
def _ask_proxy_url(*, blank_is_direct: bool) -> str | None:
_print_proxy_help()
if blank_is_direct:
_info("Leave blank for a direct connection.")
else:
_info("Leave blank to keep the current value, or type none for direct.")
while True:
value = _prompt("HTTP(S) proxy URL")
if not value:
return None if blank_is_direct else ""
if not blank_is_direct and value.lower() in ("none", "off", "direct"):
return None
if valid_proxy_url(value):
return value
_err(
"Need an HTTP or HTTPS URL like http://host:port or "
"http://user:pass@host:port — not SOCKS."
)
def _instance_telemetry() -> dict:
return {
"telemetry_enabled": 1,
"telemetry_url": DEFAULT_TELEMETRY_URL,
}
def valid_email(s: str) -> bool:
return bool(_EMAIL_RE.match(s or ""))
async def validate_token(http, token: str) -> dict | None:
try:
async with http.get(
f"{DISCORD_API}/users/@me",
headers={"Authorization": f"Bot {token}"},
) as resp:
if resp.status != 200:
return None
return await resp.json()
except aiohttp.ClientError:
return None
async def member_in_guild(http, bot_token, guild_id, user_id) -> bool:
try:
async with http.get(
f"{DISCORD_API}/guilds/{guild_id}/members/{user_id}",
headers={"Authorization": f"Bot {bot_token}"},
) as resp:
return resp.status == 200
except aiohttp.ClientError:
return False
def validate_imap(host, port, user, password) -> bool:
return validate_imap_cfg(
{"host": host, "port": port, "username": user, "password": password}
)
def imap_choice_default(email: str) -> str:
return {"gmail": "1", "outlook": "2"}.get(guess_provider(email), "3")
def _oauth_imap_cfg(name: str, email: str, tokens: dict) -> dict | None:
refresh = tokens.get("refresh_token")
if not refresh or not tokens.get("access_token"):
return None
p = provider(name)
cfg = {
"host": p["host"],
"port": p["port"],
"username": email,
"password": None,
"auth": "xoauth2",
"provider": name,
"refresh_token": refresh,
}
if not validate_imap_cfg(cfg):
return None
return {
"imap_host": p["host"],
"imap_port": p["port"],
"imap_user": email,
"imap_password": None,
"imap_auth": "xoauth2",
"imap_provider": name,
"imap_refresh_token": refresh,
}
def _paste_code_setup(name: str, security_email: str) -> dict | None:
verifier, challenge = pkce_pair()
httpd = None
redirect_uri = provider(name)["redirect_uri"]
try:
httpd, redirect_uri = start_loopback_receiver()
except OSError:
httpd = None
url = authorize_url(
name, challenge, login_hint=security_email, redirect_uri=redirect_uri
)
_info(f"Open this URL and sign in as {security_email}")
print(f" {url}")
_info(
f"After you allow access, your browser may fail to load {redirect_uri} — "
"paste the code (or the whole redirect URL) here. If the page loaded, press Enter."
)
heard = {"code": ""}
def _listen():
if httpd is None:
return
try:
heard["code"] = wait_loopback_code(httpd)
except Exception:
pass
listener = threading.Thread(target=_listen, daemon=True)
listener.start()
pasted = parse_auth_code(_prompt("Code"))
if pasted:
code = pasted
else:
listener.join(timeout=300)
code = parse_auth_code(heard["code"])
if httpd is not None:
try:
httpd.server_close()
except OSError:
pass
if not code:
_err("No authorization code found in that input.")
return None
try:
tokens = exchange_code(name, code, verifier, redirect_uri=redirect_uri)
except OAuthError as e:
_err(e.description)
return None
result = _oauth_imap_cfg(name, security_email, tokens)
if result is None:
_err(
"Sign-in succeeded but IMAP login failed — is IMAP enabled, "
"and did you sign in as that address?"
)
else:
_ok("IMAP sign-in succeeded.")
return result
def _gmail_oauth_setup(security_email: str) -> dict | None:
return _paste_code_setup("gmail", security_email)
def _outlook_oauth_setup(security_email: str) -> dict | None:
try:
start = start_device("outlook")
except OAuthError as e:
_info(
f"Device sign-in isn't available ({e.description}); using the browser-code flow."
)
return _paste_code_setup("outlook", security_email)
_info("Open this page on any device and enter the code:")
print(f" {start.get('verification_uri')}")
print(f" Code: {start.get('user_code')}")
complete = start.get("verification_uri_complete")
if complete:
print(f" Direct link: {complete}")
_wait("Waiting for you to finish signing in...")
try:
tokens = device_code_wait("outlook", start)
except OAuthError as e:
_err(e.description)
return None
result = _oauth_imap_cfg("outlook", security_email, tokens)
if result is None:
_err(
"Microsoft sign-in succeeded but IMAP login failed — did you sign in as that address?"
)
else:
_ok("IMAP sign-in succeeded.")
return result
def _password_imap_setup(security_email: str) -> dict | None:
imap_host = _prompt("IMAP host")
imap_port = int(_prompt("IMAP port", default="993"))
imap_user = _prompt("IMAP username", default=security_email)
imap_password = getpass.getpass(" IMAP password: ")
if not validate_imap(imap_host, imap_port, imap_user, imap_password):
_err("IMAP login failed — please re-check those credentials.")
return None
_ok("IMAP login succeeded.")
return {
"imap_host": imap_host,
"imap_port": imap_port,
"imap_user": imap_user,
"imap_password": imap_password,
"imap_auth": "password",
"imap_provider": None,
"imap_refresh_token": None,
}
async def _collect_mail(security_email: str) -> dict:
_info("How should the bot read security mail at that address?")
print(" 1) Gmail (sign in with Google)")
print(" 2) Outlook / Hotmail / Microsoft 365 (sign in with Microsoft)")
print(" 3) Other IMAP (username + password)")
choice = _prompt("Choice", default=imap_choice_default(security_email))
while True:
if choice == "1":
result = await asyncio.to_thread(_gmail_oauth_setup, security_email)
elif choice == "2":
result = await asyncio.to_thread(_outlook_oauth_setup, security_email)
else:
result = await asyncio.to_thread(_password_imap_setup, security_email)
if result:
return result
choice = _prompt("Choice", default=choice)
def invite_url(client_id: str) -> str:
return (
"https://discord.com/api/oauth2/authorize"
f"?client_id={client_id}"
"&scope=bot%20applications.commands"
f"&permissions={INVITE_PERMISSIONS}"
)
async def _wait_until_bot_joined(http, token, guild_id, poll_seconds=3) -> None:
while True:
try:
async with http.get(
f"{DISCORD_API}/users/@me/guilds",
headers={"Authorization": f"Bot {token}"},
) as resp:
if resp.status == 200:
guilds = await resp.json()
if any(str(g.get("id")) == str(guild_id) for g in guilds):
return
except aiohttp.ClientError:
pass
await asyncio.sleep(poll_seconds)
async def _create_secure_role(http, token, guild_id) -> str | None:
try:
async with http.post(
f"{DISCORD_API}/guilds/{guild_id}/roles",
headers={"Authorization": f"Bot {token}"},
json={"name": SECURE_ROLE_NAME, "permissions": "0"},
) as resp:
if resp.status not in (200, 201):
return None
role = await resp.json()
return str(role["id"])
except aiohttp.ClientError:
return None
async def _assign_role(http, token, guild_id, user_id, role_id) -> bool:
async with http.put(
f"{DISCORD_API}/guilds/{guild_id}/members/{user_id}/roles/{role_id}",
headers={"Authorization": f"Bot {token}"},
) as resp:
return resp.status in (200, 204)
def _print_instance_summary(row: dict) -> None:
print()
print(_paint(_BOLD, " Current instance"))
print(f" Security email {row.get('security_email') or '—'}")
host = row.get("imap_host") or "—"
port = row.get("imap_port")
location = f"{host}:{port}" if port else host
provider = row.get("imap_provider")
auth = row.get("imap_auth") or "password"
extra = f"{provider}, {auth}" if provider else auth
print(f" IMAP {location} ({extra})")
print(f" Proxy {_mask_proxy(row.get('proxy_url'))}")
async def _edit_proxy(db, existing: dict) -> None:
current = existing.get("proxy_url") or ""
_info(f"Current proxy: {_mask_proxy(current)}")
value = _ask_proxy_url(blank_is_direct=False)
if value == "":
_ok("Proxy unchanged.")
return
if value is None:
await store.set_instance_config(db, proxy_url=None)
existing["proxy_url"] = None
_ok("Proxy cleared — using direct connect.")
return
await store.set_instance_config(db, proxy_url=value)
existing["proxy_url"] = value
_ok(f"Proxy set to {_mask_proxy(value)}.")
async def _edit_mail(db, existing: dict) -> None:
current = existing.get("security_email") or ""
security_email = _prompt("Security notification email", default=current)
while not valid_email(security_email):
_err("That doesn't look like a valid email address.")
security_email = _prompt("Security notification email", default=current)
mail = await _collect_mail(security_email)
await store.set_instance_config(db, security_email=security_email, **mail)
existing["security_email"] = security_email
existing.update(mail)
_ok("Security email and IMAP updated.")
async def _edit_wizard(db, http, existing: dict) -> None:
_banner("Edit setup")
while True:
_print_instance_summary(existing)
print()
print(" 1) Add another Discord bot")
print(" 2) Change proxy")
print(" 3) Change security email / IMAP")
print(" 4) Done")
choice = _prompt("Choice", default="4")
if choice == "1":
await _setup_one_bot(db, http, True)
elif choice == "2":
await _edit_proxy(db, existing)
elif choice == "3":
await _edit_mail(db, existing)
elif choice == "4":
_ok("Setup complete.")
return
else:
_err("Choose 1, 2, 3, or 4.")
async def _setup_one_bot(db, http, instance_configured: bool) -> bool:
total = 2 if instance_configured else 4
_step(1, total, "Discord bot")
token = _prompt("Discord bot token")
bot_user = await validate_token(http, token)
while bot_user is None:
_err("Could not validate that token against Discord's API — try again.")
token = _prompt("Discord bot token")
bot_user = await validate_token(http, token)
client_id = str(bot_user["id"])
label = bot_user.get("username") or client_id
_ok(f"Authenticated as {label}")
_info("Invite this bot with the least-privilege link:")
print(f" {invite_url(client_id)}")
guild_id = _prompt("Guild (server) ID the bot was invited to")
_wait("Waiting for the bot to join that server...")
await _wait_until_bot_joined(http, token, guild_id)
_ok("Bot detected in the guild.")
_step(2, total, "Secure Access role")
role_id = await _create_secure_role(http, token, guild_id)
while role_id is None:
_err(
f"Couldn't create the '{SECURE_ROLE_NAME}' role — check the bot has "
"the Manage Roles permission in that server."
)
_pause()
role_id = await _create_secure_role(http, token, guild_id)
_ok(f"Created '{SECURE_ROLE_NAME}' role.")
user_id = _prompt("Discord user ID to assign the Secure Access role to")
while not await member_in_guild(http, token, guild_id, user_id):
_err("That user isn't a member of the guild yet — try again.")
user_id = _prompt("Discord user ID to assign the Secure Access role to")
while not await _assign_role(http, token, guild_id, user_id, role_id):
_err(
"Couldn't assign the role — check the bot's own role is positioned "
"above 'Secure Access' in Server Settings -> Roles."
)
_pause()
_ok("Role assigned.")
if instance_configured:
_info("Reusing this instance's security email, IMAP, and proxy settings.")
else:
_step(3, total, "Security email")
security_email = _prompt("Security notification email")
while not valid_email(security_email):
_err("That doesn't look like a valid email address.")
security_email = _prompt("Security notification email")
mail = await _collect_mail(security_email)
_step(4, total, "Outbound proxy")
_info("Used for Microsoft login and IMAP traffic.")
proxy_url = _ask_proxy_url(blank_is_direct=True)
await store.set_instance_config(
db,
security_email=security_email,
proxy_url=proxy_url,
**_instance_telemetry(),
**mail,
)
instance_configured = True
_ok("Instance settings saved.")
bot_id = await store.add_bot(db, token, label)
await store.set_guild_config(db, int(guild_id), bot_id, secure_role_id=int(role_id))
_ok(f"Bot '{label}' configured and saved.")
return instance_configured
async def run_wizard(db, http) -> None:
existing = await store.get_instance_config(db)
instance_configured = bool(existing and existing.get("security_email"))
if instance_configured:
await _edit_wizard(db, http, dict(existing))
return
_banner("Setup wizard")
instance_configured = await _setup_one_bot(db, http, False)
while _yes("Set up another bot?"):
instance_configured = await _setup_one_bot(db, http, instance_configured)
print()
_ok("Setup complete. Start the bot with: python -m bot")
async def _main() -> None:
db = await store.connect(resolve_db_path(os.environ.get("OAS_DB_PATH")))
try:
async with aiohttp.ClientSession() as http:
await run_wizard(db, http)
finally:
await db.close()
if __name__ == "__main__":
asyncio.run(_main())