diff --git a/app/repositories/portfolio_repository.py b/app/repositories/portfolio_repository.py index 71b8ec6..02cd203 100644 --- a/app/repositories/portfolio_repository.py +++ b/app/repositories/portfolio_repository.py @@ -5,7 +5,7 @@ from __future__ import annotations from datetime import datetime from typing import Protocol -from app.domain.models import AuthToken, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, TokenPrincipal, User +from app.domain.models import AuthToken, BrokerCredential, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, TokenPrincipal, User class PortfolioRepository(Protocol): @@ -42,6 +42,19 @@ class PortfolioRepository(Protocol): def update_user_last_login(self, user_id: str) -> None: ... + def upsert_broker_credential( + self, + user_id: str, + broker: str, + *, + account_ref: str, + encrypted_payload: str, + ) -> BrokerCredential: ... + + def list_broker_credentials(self, user_id: str) -> list[BrokerCredential]: ... + + def delete_broker_credential(self, user_id: str, broker: str) -> bool: ... + def create_auth_token( self, user_id: str, diff --git a/app/repositories/postgres.py b/app/repositories/postgres.py index 2ac6f0e..4e7f3fa 100644 --- a/app/repositories/postgres.py +++ b/app/repositories/postgres.py @@ -11,7 +11,7 @@ from typing import Any import psycopg from psycopg import sql -from app.domain.models import AuthToken, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, TokenPrincipal, User +from app.domain.models import AuthToken, BrokerCredential, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, TokenPrincipal, User _SELECT_USER = """ SELECT @@ -331,6 +331,19 @@ class PostgresPortfolioRepository: daemon_id=row[11] if len(row) > 11 else None, ) + @staticmethod + def _row_to_broker_credential(row: Any) -> BrokerCredential: + if row is None: + raise ValueError("브로커 인증 정보를 찾을 수 없습니다") + return BrokerCredential( + user_id=row[0], + broker=row[1], + account_ref=row[2], + encrypted_payload=row[3], + credential_version=int(row[4]) if len(row) > 4 else 1, + last_rotated_at=row[5] if len(row) > 5 else None, + ) + @staticmethod def _row_to_portfolio(row: Any) -> Portfolio: if row is None: @@ -552,6 +565,60 @@ class PostgresPortfolioRepository: (user_id,), ) + def upsert_broker_credential( + self, + user_id: str, + broker: str, + *, + account_ref: str, + encrypted_payload: str, + ) -> BrokerCredential: + self.ensure_schema() + with self._connect().cursor() as cur: + cur.execute("SELECT 1 FROM users WHERE id = %s", (user_id,)) + if cur.fetchone() is None: + raise ValueError("사용자를 찾을 수 없습니다") + cur.execute( + """ + INSERT INTO broker_credentials (user_id, broker, account_ref, encrypted_payload) + VALUES (%s, %s, %s, %s) + ON CONFLICT (user_id, broker) DO UPDATE + SET account_ref = EXCLUDED.account_ref, + encrypted_payload = EXCLUDED.encrypted_payload, + credential_version = broker_credentials.credential_version + 1, + last_rotated_at = NOW() + RETURNING user_id, broker, account_ref, encrypted_payload, credential_version, last_rotated_at + """, + (user_id, broker, account_ref, encrypted_payload), + ) + row = cur.fetchone() + return self._row_to_broker_credential(row) + + def list_broker_credentials(self, user_id: str) -> list[BrokerCredential]: + self.ensure_schema() + with self._connect().cursor() as cur: + cur.execute( + """ + SELECT user_id, broker, account_ref, encrypted_payload, credential_version, last_rotated_at + FROM broker_credentials + WHERE user_id = %s + ORDER BY broker ASC + """, + (user_id,), + ) + rows = cur.fetchall() + return [self._row_to_broker_credential(row) for row in rows] + + def delete_broker_credential(self, user_id: str, broker: str) -> bool: + self.ensure_schema() + with self._connect().cursor() as cur: + cur.execute( + "DELETE FROM broker_credentials WHERE user_id = %s AND broker = %s RETURNING user_id", + (user_id, broker), + ) + row = cur.fetchone() + return row is not None + def create_auth_token( self, user_id: str, diff --git a/tests/test_cli.py b/tests/test_cli.py index dd466e2..768b43c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -10,7 +10,7 @@ from typer.testing import CliRunner from app.cli.main import create_app from app.core.config import Settings from app.core.security import hash_token, token_prefix -from app.domain.models import AuthToken, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, User +from app.domain.models import AuthToken, BrokerCredential, Daemon, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, User class FakeRepo: @@ -26,6 +26,7 @@ class FakeRepo: self._comments: dict[str, PortfolioComment] = {} self._tags: dict[str, PortfolioTag] = {} self._auth_tokens: dict[str, AuthToken] = {} + self._broker_credentials: dict[tuple[str, str], BrokerCredential] = {} self._token_seq = 1 self.closed = False @@ -167,6 +168,41 @@ class FakeRepo: daemon.updated_at = datetime.now(timezone.utc) return daemon + def upsert_broker_credential( + self, + user_id: str, + broker: str, + *, + account_ref: str, + encrypted_payload: str, + ) -> BrokerCredential: + self.get_user(user_id) + key = (user_id, broker) + existing = self._broker_credentials.get(key) + credential = BrokerCredential( + user_id=user_id, + broker=broker, + account_ref=account_ref, + encrypted_payload=encrypted_payload, + credential_version=(existing.credential_version + 1) if existing is not None else 1, + last_rotated_at=datetime.now(timezone.utc), + ) + self._broker_credentials[key] = credential + return credential + + def list_broker_credentials(self, user_id: str) -> list[BrokerCredential]: + return sorted( + [credential for (owner_user_id, _), credential in self._broker_credentials.items() if owner_user_id == user_id], + key=lambda credential: credential.broker, + ) + + def delete_broker_credential(self, user_id: str, broker: str) -> bool: + key = (user_id, broker) + if key not in self._broker_credentials: + return False + del self._broker_credentials[key] + return True + def create_auth_token( self, user_id: str, @@ -481,7 +517,11 @@ class FakeRepo: def _settings() -> Settings: - return Settings(db_enabled=True, allowed_origins_raw="http://localhost:3000") + return Settings( + db_enabled=True, + allowed_origins_raw="http://localhost:3000", + credential_encryption_key="7ApeKhDwGrxiPdlfNbWw2pylyU-XvTk4v3CiRDnNdV0=", + ) def build_runner(repo: FakeRepo) -> CliRunner: @@ -928,6 +968,63 @@ def test_cli_daemon_create_and_list_render_entries() -> None: assert "Alpha Bot" in list_result.stdout +def test_cli_credential_set_list_and_remove_flow() -> None: + repo = FakeRepo(_portfolio()) + user = repo.create_user("alice", permission="read") + runner = build_runner(repo) + + set_result = runner.invoke( + runner.app, + [ + "--format", + "json", + "credential", + "set", + "--user", + "alice", + "--broker", + "kis", + "--account-ref", + "123-45", + "--app-key", + "key-1", + "--app-secret", + "secret-1", + "--force", + ], + ) + list_result = runner.invoke( + runner.app, + ["--format", "json", "credential", "list", "--user", "alice"], + ) + remove_result = runner.invoke( + runner.app, + ["--format", "json", "credential", "remove", "--user", "alice", "--broker", "kis", "--force"], + ) + + assert set_result.exit_code == 0 + set_payload = json.loads(set_result.stdout) + assert set_payload["user_id"] == user.id + assert set_payload["broker"] == "kis" + assert set_payload["credential_version"] == 1 + + assert list_result.exit_code == 0 + list_payload = json.loads(list_result.stdout) + assert list_payload == [ + { + "broker": "kis", + "account_ref": "123-45", + "credential_version": 1, + "last_rotated_at": list_payload[0]["last_rotated_at"], + } + ] + + assert remove_result.exit_code == 0 + remove_payload = json.loads(remove_result.stdout) + assert remove_payload["deleted"] is True + assert repo._broker_credentials == {} + + def test_cli_service_token_issue_supports_json() -> None: repo = FakeRepo(_portfolio()) user = repo.create_user("alice", permission="read")