FluxIP/app/accounts/repository.py

643 lines
25 KiB
Python

from __future__ import annotations
import json
import sqlite3
import uuid
from collections.abc import Iterable, Mapping
from dataclasses import asdict, dataclass
from typing import Any, Literal
from app.core.security import EncryptedValue, SecretCipher
from app.core.time import to_iso
from app.database.database import Database
AccountProvider = Literal["aws", "cloudflare"]
AWS_SECRET_NAMES = (
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
)
CLOUDFLARE_SECRET_NAMES = ("cloudflare_api_token",)
SECRET_NAMES = (*AWS_SECRET_NAMES, *CLOUDFLARE_SECRET_NAMES)
ACCOUNT_COLUMNS = """
id, provider, name, use_default_aws_credentials, config_version,
created_at, updated_at, archived_at
"""
@dataclass(slots=True)
class CredentialAccountRecord:
id: str
provider: AccountProvider
name: str
use_default_aws_credentials: bool | None
config_version: int
created_at: str
updated_at: str
archived_at: str | None
def to_dict(self) -> dict[str, Any]:
return asdict(self)
class AccountRepository:
def __init__(self, database: Database, cipher: SecretCipher) -> None:
self.database = database
self.cipher = cipher
def list_accounts(
self,
provider: AccountProvider | None = None,
*,
include_archived: bool = False,
) -> list[CredentialAccountRecord]:
clauses: list[str] = []
parameters: list[object] = []
if provider is not None:
clauses.append("provider = ?")
parameters.append(provider)
if not include_archived:
clauses.append("archived_at IS NULL")
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
with self.database.connect() as connection:
rows = connection.execute(
f"SELECT {ACCOUNT_COLUMNS} FROM credential_accounts {where} "
"ORDER BY provider, name COLLATE NOCASE, id",
parameters,
).fetchall()
return [self._account_from_row(row) for row in rows]
def get_account(
self,
account_id: str,
*,
include_archived: bool = False,
) -> CredentialAccountRecord | None:
archived_clause = "" if include_archived else " AND archived_at IS NULL"
with self.database.connect() as connection:
row = connection.execute(
f"SELECT {ACCOUNT_COLUMNS} FROM credential_accounts WHERE id = ?{archived_clause}",
(account_id,),
).fetchone()
return self._account_from_row(row) if row else None
def create_account(
self,
*,
provider: AccountProvider,
name: str,
use_default_aws_credentials: bool | None,
secret_values: Mapping[str, str],
) -> CredentialAccountRecord:
account_id = str(uuid.uuid4())
now = to_iso()
self._validate_secret_names(provider, secret_values)
with self.database.connect() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
self._ensure_writes_allowed(connection)
connection.execute(
"""
INSERT INTO credential_accounts(
id, provider, name, use_default_aws_credentials,
config_version, created_at, updated_at
) VALUES (?, ?, ?, ?, 1, ?, ?)
""",
(
account_id,
provider,
name.strip(),
self._database_bool(use_default_aws_credentials),
now,
now,
),
)
self._upsert_secrets(connection, account_id, secret_values, now)
connection.commit()
except sqlite3.IntegrityError as exc:
connection.rollback()
raise RuntimeError(self._integrity_code(exc)) from exc
except Exception:
connection.rollback()
raise
created = self.get_account(account_id)
if created is None: # pragma: no cover - insert and read share the same database
raise RuntimeError("ACCOUNT_CREATION_FAILED")
return created
def import_legacy_accounts(
self,
*,
aws_use_default: bool | None,
aws_secret_values: Mapping[str, str],
cloudflare_api_token: str | None,
) -> dict[AccountProvider, str]:
"""Import the former global credentials exactly once during startup."""
imported: dict[AccountProvider, str] = {}
now = to_iso()
with self.database.connect() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
if aws_use_default is not None:
aws_id, aws_can_import = self._legacy_account_target(connection, "aws")
if aws_id is None and aws_can_import:
aws_id = "legacy-aws"
connection.execute(
"""
INSERT INTO credential_accounts(
id, provider, name, use_default_aws_credentials,
config_version, created_at, updated_at
) VALUES (?, 'aws', ?, ?, 1, ?, ?)
""",
(
aws_id,
"旧版全局 AWS 账号",
int(aws_use_default),
now,
now,
),
)
self._upsert_secrets(
connection,
aws_id,
aws_secret_values,
now,
)
if aws_id == "legacy-aws" and aws_can_import:
imported["aws"] = aws_id
connection.execute(
"""
UPDATE managed_instances SET aws_account_id = ?
WHERE aws_account_id IS NULL
""",
(aws_id,),
)
connection.execute(
"""
UPDATE fleet_run_items SET aws_account_id = ?
WHERE aws_account_id IS NULL
""",
(aws_id,),
)
if cloudflare_api_token:
cloudflare_id, cloudflare_can_import = self._legacy_account_target(
connection,
"cloudflare",
)
if cloudflare_id is None and cloudflare_can_import:
cloudflare_id = "legacy-cloudflare"
connection.execute(
"""
INSERT INTO credential_accounts(
id, provider, name, use_default_aws_credentials,
config_version, created_at, updated_at
) VALUES (?, 'cloudflare', ?, NULL, 1, ?, ?)
""",
(
cloudflare_id,
"旧版全局 Cloudflare 账号",
now,
now,
),
)
self._upsert_secrets(
connection,
cloudflare_id,
{"cloudflare_api_token": cloudflare_api_token},
now,
)
if cloudflare_id == "legacy-cloudflare" and cloudflare_can_import:
imported["cloudflare"] = cloudflare_id
connection.execute(
"""
UPDATE managed_instances SET cloudflare_account_id = ?
WHERE cloudflare_account_id IS NULL
""",
(cloudflare_id,),
)
connection.execute(
"""
UPDATE fleet_run_items SET cloudflare_account_id = ?
WHERE cloudflare_account_id IS NULL
""",
(cloudflare_id,),
)
connection.commit()
except sqlite3.IntegrityError as exc:
connection.rollback()
raise RuntimeError(self._integrity_code(exc)) from exc
except Exception:
connection.rollback()
raise
return imported
def update_account(
self,
account_id: str,
*,
expected_version: int,
values: Mapping[str, object],
secret_values: Mapping[str, str],
delete_secret_names: Iterable[str] = (),
) -> CredentialAccountRecord:
delete_names = tuple(dict.fromkeys(delete_secret_names))
with self.database.connect() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
current = self._get_account_row(connection, account_id)
if current is None:
raise RuntimeError("ACCOUNT_NOT_FOUND")
if int(current["config_version"]) != expected_version:
raise RuntimeError("ACCOUNT_CONFIG_VERSION_CONFLICT")
provider = str(current["provider"])
self._validate_secret_names(provider, secret_values)
self._validate_secret_names(provider, delete_names)
recovery_target = self._ensure_account_update_allowed(
connection,
account_id,
current,
values,
secret_values,
)
updates: dict[str, object] = {}
if "name" in values:
updates["name"] = str(values["name"]).strip()
if "use_default_aws_credentials" in values:
updates["use_default_aws_credentials"] = self._database_bool(
values["use_default_aws_credentials"]
)
now = to_iso()
if updates:
assignments = ", ".join(f"{field} = ?" for field in updates)
connection.execute(
f"""
UPDATE credential_accounts
SET {assignments}, config_version = config_version + 1,
updated_at = ?
WHERE id = ? AND archived_at IS NULL
""",
(*updates.values(), now, account_id),
)
else:
connection.execute(
"""
UPDATE credential_accounts
SET config_version = config_version + 1, updated_at = ?
WHERE id = ? AND archived_at IS NULL
""",
(now, account_id),
)
if delete_names:
placeholders = ", ".join("?" for _ in delete_names)
connection.execute(
f"DELETE FROM account_secrets "
f"WHERE account_id = ? AND name IN ({placeholders})",
(account_id, *delete_names),
)
self._upsert_secrets(connection, account_id, secret_values, now)
if recovery_target is not None:
run_id, item_id = recovery_target
connection.execute(
"""
INSERT INTO fleet_events(
run_id, item_id, occurred_at, stage, level,
message, details_json
) VALUES (?, ?, ?, 'credentials', 'info', ?, ?)
""",
(
run_id,
item_id,
now,
"管理员已更新当前账号的恢复凭据",
json.dumps(
{
"account_id": account_id,
"updated_credentials": sorted(
set(secret_values) | set(delete_names)
),
},
ensure_ascii=False,
),
),
)
connection.commit()
except sqlite3.IntegrityError as exc:
connection.rollback()
raise RuntimeError(self._integrity_code(exc)) from exc
except Exception:
connection.rollback()
raise
updated = self.get_account(account_id)
if updated is None: # pragma: no cover
raise RuntimeError("ACCOUNT_NOT_FOUND")
return updated
def archive_account(self, account_id: str) -> CredentialAccountRecord:
now = to_iso()
with self.database.connect() as connection:
try:
connection.execute("BEGIN IMMEDIATE")
self._ensure_writes_allowed(connection)
if self._get_account_row(connection, account_id) is None:
raise RuntimeError("ACCOUNT_NOT_FOUND")
if connection.execute(
"""
SELECT 1 FROM managed_instances
WHERE archived_at IS NULL
AND (aws_account_id = ? OR cloudflare_account_id = ?)
LIMIT 1
""",
(account_id, account_id),
).fetchone():
raise RuntimeError("ACCOUNT_IN_USE")
connection.execute(
"DELETE FROM account_secrets WHERE account_id = ?",
(account_id,),
)
connection.execute(
"""
UPDATE credential_accounts
SET config_version = config_version + 1,
updated_at = ?, archived_at = ?
WHERE id = ? AND archived_at IS NULL
""",
(now, now, account_id),
)
connection.commit()
except Exception:
connection.rollback()
raise
archived = self.get_account(account_id, include_archived=True)
if archived is None: # pragma: no cover
raise RuntimeError("ACCOUNT_NOT_FOUND")
return archived
def secret_flags(self, account_id: str) -> dict[str, bool]:
return self.secret_flags_many([account_id]).get(
account_id,
{name: False for name in SECRET_NAMES},
)
def secret_flags_many(self, account_ids: Iterable[str]) -> dict[str, dict[str, bool]]:
ids = list(dict.fromkeys(account_ids))
result = {account_id: {name: False for name in SECRET_NAMES} for account_id in ids}
if not ids:
return result
placeholders = ", ".join("?" for _ in ids)
with self.database.connect() as connection:
rows = connection.execute(
f"SELECT account_id, name FROM account_secrets "
f"WHERE account_id IN ({placeholders})",
ids,
).fetchall()
for row in rows:
account_id = str(row["account_id"])
name = str(row["name"])
if account_id in result and name in result[account_id]:
result[account_id][name] = True
return result
def resolve_aws(
self,
account_id: str,
) -> tuple[bool, dict[str, str | None]]:
account = self._require_provider(account_id, "aws")
secrets = {name: self._get_secret(account.id, name) for name in AWS_SECRET_NAMES}
use_default = bool(account.use_default_aws_credentials)
if not use_default and not (
secrets["aws_access_key_id"] and secrets["aws_secret_access_key"]
):
raise RuntimeError("ACCOUNT_CREDENTIALS_INCOMPLETE")
return use_default, secrets
def resolve_cloudflare(self, account_id: str) -> str:
account = self._require_provider(account_id, "cloudflare")
token = self._get_secret(account.id, "cloudflare_api_token")
if not token:
raise RuntimeError("ACCOUNT_CREDENTIALS_INCOMPLETE")
return token
def _require_provider(
self,
account_id: str,
provider: AccountProvider,
) -> CredentialAccountRecord:
account = self.get_account(account_id)
if account is None:
raise RuntimeError("ACCOUNT_NOT_FOUND")
if account.provider != provider:
raise RuntimeError("ACCOUNT_PROVIDER_MISMATCH")
return account
def _get_secret(self, account_id: str, name: str) -> str | None:
with self.database.connect() as connection:
row = connection.execute(
"""
SELECT key_version, nonce, ciphertext
FROM account_secrets WHERE account_id = ? AND name = ?
""",
(account_id, name),
).fetchone()
if row is None:
return None
return self.cipher.decrypt(
self._secret_aad(account_id, name),
EncryptedValue(**dict(row)),
)
@staticmethod
def _legacy_account_target(
connection: sqlite3.Connection,
provider: AccountProvider,
) -> tuple[str | None, bool]:
legacy_id = f"legacy-{provider}"
legacy = connection.execute(
"SELECT archived_at FROM credential_accounts WHERE id = ?",
(legacy_id,),
).fetchone()
if legacy is not None:
return (
legacy_id if legacy["archived_at"] is None else None,
legacy["archived_at"] is None,
)
active = connection.execute(
"""
SELECT id FROM credential_accounts
WHERE provider = ? AND archived_at IS NULL
ORDER BY created_at, id
LIMIT 1
""",
(provider,),
).fetchone()
if active is not None:
return str(active["id"]), False
return None, True
def _upsert_secrets(
self,
connection: sqlite3.Connection,
account_id: str,
secret_values: Mapping[str, str],
now: str,
) -> None:
for name, plaintext in secret_values.items():
encrypted = self.cipher.encrypt(
self._secret_aad(account_id, name),
plaintext.strip(),
)
connection.execute(
"""
INSERT INTO account_secrets(
account_id, name, key_version, nonce, ciphertext,
created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(account_id, name) DO UPDATE SET
key_version = excluded.key_version,
nonce = excluded.nonce,
ciphertext = excluded.ciphertext,
updated_at = excluded.updated_at
""",
(
account_id,
name,
encrypted.key_version,
encrypted.nonce,
encrypted.ciphertext,
now,
now,
),
)
@staticmethod
def _secret_aad(account_id: str, name: str) -> str:
return f"credential-account:{account_id}:{name}"
@staticmethod
def _ensure_writes_allowed(connection: sqlite3.Connection) -> None:
now = to_iso()
connection.execute(
"""
DELETE FROM fleet_operation_locks
WHERE kind = 'dns_sync' AND lease_until IS NOT NULL AND lease_until < ?
""",
(now,),
)
if connection.execute("SELECT 1 FROM fleet_runs WHERE active_slot = 1 LIMIT 1").fetchone():
raise RuntimeError("FLEET_RUN_ACTIVE")
if connection.execute(
"""
SELECT 1 FROM fleet_operation_locks
WHERE lease_until IS NULL OR lease_until >= ?
LIMIT 1
""",
(now,),
).fetchone():
raise RuntimeError("FLEET_RUN_ACTIVE")
@staticmethod
def _ensure_account_update_allowed(
connection: sqlite3.Connection,
account_id: str,
current: sqlite3.Row,
values: Mapping[str, object],
secret_values: Mapping[str, str],
) -> tuple[str, str | None] | None:
now = to_iso()
connection.execute(
"""
DELETE FROM fleet_operation_locks
WHERE kind = 'dns_sync' AND lease_until IS NOT NULL AND lease_until < ?
""",
(now,),
)
active = connection.execute(
"""
SELECT id, status, current_item_id
FROM fleet_runs WHERE active_slot = 1
"""
).fetchone()
if active is None:
if connection.execute("SELECT 1 FROM fleet_operation_locks LIMIT 1").fetchone():
raise RuntimeError("FLEET_RUN_ACTIVE")
return None
if active["status"] not in {"needs_attention", "cleanup_pending"}:
raise RuntimeError("FLEET_RUN_ACTIVE")
if connection.execute(
"""
SELECT 1 FROM fleet_operation_locks
WHERE NOT (kind = 'rotation' AND owner_id = ?)
LIMIT 1
""",
(active["id"],),
).fetchone():
raise RuntimeError("FLEET_RUN_ACTIVE")
item = connection.execute(
"""
SELECT 1 FROM fleet_run_items
WHERE id = ? AND (aws_account_id = ? OR cloudflare_account_id = ?)
""",
(active["current_item_id"], account_id, account_id),
).fetchone()
if item is None:
raise RuntimeError("FLEET_RUN_ACTIVE")
name_changed = "name" in values and str(values["name"]).strip() != str(current["name"])
mode_changed = (
"use_default_aws_credentials" in values
and AccountRepository._database_bool(values["use_default_aws_credentials"])
!= current["use_default_aws_credentials"]
)
if name_changed or mode_changed or not secret_values:
raise RuntimeError("FLEET_RUN_ACTIVE")
return str(active["id"]), (
str(active["current_item_id"]) if active["current_item_id"] else None
)
@staticmethod
def _validate_secret_names(
provider: str,
values: Mapping[str, object] | Iterable[str],
) -> None:
names = set(values)
allowed = set(AWS_SECRET_NAMES if provider == "aws" else CLOUDFLARE_SECRET_NAMES)
if not names <= allowed:
raise RuntimeError("ACCOUNT_PROVIDER_MISMATCH")
@staticmethod
def _database_bool(value: object) -> int | None:
if value is None:
return None
return int(bool(value))
@staticmethod
def _get_account_row(
connection: sqlite3.Connection,
account_id: str,
) -> sqlite3.Row | None:
return connection.execute(
f"SELECT {ACCOUNT_COLUMNS} FROM credential_accounts "
"WHERE id = ? AND archived_at IS NULL",
(account_id,),
).fetchone()
@staticmethod
def _account_from_row(row: sqlite3.Row) -> CredentialAccountRecord:
values = dict(row)
if values["use_default_aws_credentials"] is not None:
values["use_default_aws_credentials"] = bool(values["use_default_aws_credentials"])
return CredentialAccountRecord(**values)
@staticmethod
def _integrity_code(exc: sqlite3.IntegrityError) -> str:
message = str(exc)
if "credential_accounts.provider" in message and "credential_accounts.name" in message:
return "ACCOUNT_NAME_CONFLICT"
if "credential_accounts.id" in message:
return "ACCOUNT_ID_CONFLICT"
if "CHECK constraint failed" in message or "NOT NULL constraint failed" in message:
return "INVALID_ACCOUNT_DATA"
return "ACCOUNT_CONSTRAINT_CONFLICT"