From 14df805209639ccd386c5f7a9bc5856a766573bb Mon Sep 17 00:00:00 2001 From: Noloquideus Date: Mon, 8 Jun 2026 10:25:51 +0300 Subject: [PATCH] feat: add more options for wallets --- src/application/commands/__init__.py | 8 ++ .../commands/create_organization_wallets.py | 10 ++- .../commands/organization_wallet_commands.py | 63 +++++++++++++ .../domain/entities/organization.py | 6 ++ src/infrastructure/config/settings.py | 2 +- src/infrastructure/crypto/wallet_crypto.py | 89 +++++++++++++++++++ src/presentation/dependencies/commands.py | 24 +++++ src/presentation/routing/organizations.py | 61 ++++++++++++- src/presentation/schemas/mappers.py | 29 ++++++ src/presentation/schemas/organization.py | 16 ++++ 10 files changed, 300 insertions(+), 8 deletions(-) create mode 100644 src/application/commands/organization_wallet_commands.py diff --git a/src/application/commands/__init__.py b/src/application/commands/__init__.py index a450151..306fb05 100644 --- a/src/application/commands/__init__.py +++ b/src/application/commands/__init__.py @@ -14,6 +14,11 @@ from src.application.commands.organization_document_commands import ( ListOrganizationDocumentsCommand, GetOrganizationDocumentCommand, ) +from src.application.commands.organization_wallet_commands import ( + GetOrganizationMnemonicCommand, + GetOrganizationSecretKeysCommand, + ListOrganizationWalletsCommand, +) from src.application.commands.purchase_request_commands import ( ListPurchaseRequestsCommand, GetPurchaseRequestCommand, @@ -38,4 +43,7 @@ __all__ = [ 'SetPurchaseRequestQuoteCommand', 'ListOrganizationDocumentsCommand', 'GetOrganizationDocumentCommand', + 'ListOrganizationWalletsCommand', + 'GetOrganizationMnemonicCommand', + 'GetOrganizationSecretKeysCommand', ] diff --git a/src/application/commands/create_organization_wallets.py b/src/application/commands/create_organization_wallets.py index 356c998..d338713 100644 --- a/src/application/commands/create_organization_wallets.py +++ b/src/application/commands/create_organization_wallets.py @@ -4,7 +4,10 @@ from ulid import ULID from src.application.abstractions import IUnitOfWork from src.application.contracts import ILogger -from src.application.domain.entities.organization import OrganizationWalletEntity +from src.application.domain.entities.organization import ( + CreateOrganizationWalletsResult, + OrganizationWalletEntity, +) from src.application.domain.exceptions import ApplicationException from src.infrastructure.crypto.wallet_crypto import ( ALL_CHAINS, @@ -22,7 +25,7 @@ class CreateOrganizationWalletsCommand: self._logger = logger @transactional - async def __call__(self, *, organization_id: str) -> list[OrganizationWalletEntity]: + async def __call__(self, *, organization_id: str) -> CreateOrganizationWalletsResult: if not is_crypto_ready(): raise ApplicationException(status_code=503, message='Crypto service not ready') @@ -35,6 +38,7 @@ class CreateOrganizationWalletsCommand: mnemonic = generate_mnemonic() derived = derive_all_addresses(mnemonic) + plaintext_mnemonic = mnemonic blob = encrypt_mnemonic(mnemonic) mnemonic = '' @@ -53,4 +57,4 @@ class CreateOrganizationWalletsCommand: ] saved = await self._unit_of_work.organization_wallet_repository.create_many(wallets) self._logger.info(f'Wallets created for organization_id={organization_id} chains={len(saved)}') - return saved + return CreateOrganizationWalletsResult(wallets=saved, mnemonic=plaintext_mnemonic) diff --git a/src/application/commands/organization_wallet_commands.py b/src/application/commands/organization_wallet_commands.py new file mode 100644 index 0000000..4283b67 --- /dev/null +++ b/src/application/commands/organization_wallet_commands.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from src.application.abstractions import IUnitOfWork +from src.application.contracts import ILogger +from src.application.domain.entities.organization import LegalEntityEntity, OrganizationWalletEntity +from src.application.domain.exceptions import ApplicationException +from src.infrastructure.crypto.wallet_crypto import ( + DerivedSecretKey, + decrypt_mnemonic, + derive_all_private_keys, + is_crypto_ready, +) +from src.infrastructure.database.decorators import transactional + + +class ListOrganizationWalletsCommand: + def __init__(self, unit_of_work: IUnitOfWork, logger: ILogger): + self._unit_of_work = unit_of_work + self._logger = logger + + @transactional + async def __call__(self, organization_id: str) -> list[OrganizationWalletEntity]: + await self._unit_of_work.legal_entity_repository.get_by_id(organization_id) + return await self._unit_of_work.organization_wallet_repository.list_by_organization(organization_id) + + +class GetOrganizationMnemonicCommand: + def __init__(self, unit_of_work: IUnitOfWork, logger: ILogger): + self._unit_of_work = unit_of_work + self._logger = logger + + @transactional + async def __call__(self, organization_id: str) -> str: + org = await self._require_wallets_org(organization_id) + if not is_crypto_ready(): + raise ApplicationException(status_code=503, message='Crypto service not ready') + return decrypt_mnemonic(org.encrypted_mnemonic or '') + + async def _require_wallets_org(self, organization_id: str) -> LegalEntityEntity: + org = await self._unit_of_work.legal_entity_repository.get_by_id(organization_id) + if not org.encrypted_mnemonic: + raise ApplicationException(status_code=404, message='Wallets not created for organization') + return org + + +class GetOrganizationSecretKeysCommand: + def __init__(self, unit_of_work: IUnitOfWork, logger: ILogger): + self._unit_of_work = unit_of_work + self._logger = logger + + @transactional + async def __call__(self, organization_id: str) -> list[DerivedSecretKey]: + org = await self._require_wallets_org(organization_id) + if not is_crypto_ready(): + raise ApplicationException(status_code=503, message='Crypto service not ready') + mnemonic = decrypt_mnemonic(org.encrypted_mnemonic or '') + return derive_all_private_keys(mnemonic) + + async def _require_wallets_org(self, organization_id: str) -> LegalEntityEntity: + org = await self._unit_of_work.legal_entity_repository.get_by_id(organization_id) + if not org.encrypted_mnemonic: + raise ApplicationException(status_code=404, message='Wallets not created for organization') + return org diff --git a/src/application/domain/entities/organization.py b/src/application/domain/entities/organization.py index 4436af8..7d12e72 100644 --- a/src/application/domain/entities/organization.py +++ b/src/application/domain/entities/organization.py @@ -39,6 +39,12 @@ class OrganizationWalletEntity: created_at: datetime | None = None +@dataclass +class CreateOrganizationWalletsResult: + wallets: list[OrganizationWalletEntity] + mnemonic: str + + @dataclass class OrganizationDocumentEntity: id: str diff --git a/src/infrastructure/config/settings.py b/src/infrastructure/config/settings.py index 7d343cf..c3b2cad 100644 --- a/src/infrastructure/config/settings.py +++ b/src/infrastructure/config/settings.py @@ -81,7 +81,7 @@ class Settings(BaseSettings): RATE_LIMIT_WINDOW: int = 60 LOG_LEVEL: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'] = 'INFO' - LOG_FORMAT: Literal['JSON', 'TEXT'] = 'TEXT' + LOG_FORMAT: Literal['JSON', 'TEXT'] = 'JSON' @field_validator('VAULT_ADDR', mode='before') @classmethod diff --git a/src/infrastructure/crypto/wallet_crypto.py b/src/infrastructure/crypto/wallet_crypto.py index da47c2d..23d23cf 100644 --- a/src/infrastructure/crypto/wallet_crypto.py +++ b/src/infrastructure/crypto/wallet_crypto.py @@ -43,6 +43,14 @@ class DerivedWallet: derivation_path: str +@dataclass(frozen=True) +class DerivedSecretKey: + chain: str + address: str + derivation_path: str + private_key: str + + class CryptoNotReadyError(RuntimeError): pass @@ -151,6 +159,87 @@ def derive_all_addresses(mnemonic: str) -> list[DerivedWallet]: ] +def derive_all_private_keys(mnemonic: str) -> list[DerivedSecretKey]: + seed_bytes = Bip39SeedGenerator(mnemonic).Generate() + + eth_ctx = Bip44.FromSeed(seed_bytes, Bip44Coins.ETHEREUM) + eth_node = ( + eth_ctx.Purpose() + .Coin() + .Account(0) + .Change(Bip44Changes.CHAIN_EXT) + .AddressIndex(0) + ) + eth_addr = eth_node.PublicKey().ToAddress() + eth_pk = eth_node.PrivateKey().Raw().ToHex() + + btc_ctx = Bip84.FromSeed(seed_bytes, Bip84Coins.BITCOIN) + btc_node = ( + btc_ctx.Purpose() + .Coin() + .Account(0) + .Change(Bip44Changes.CHAIN_EXT) + .AddressIndex(0) + ) + btc_addr = btc_node.PublicKey().ToAddress() + btc_pk = btc_node.PrivateKey().Raw().ToHex() + + trx_ctx = Bip44.FromSeed(seed_bytes, Bip44Coins.TRON) + trx_node = ( + trx_ctx.Purpose() + .Coin() + .Account(0) + .Change(Bip44Changes.CHAIN_EXT) + .AddressIndex(0) + ) + trx_addr = trx_node.PublicKey().ToAddress() + trx_pk = trx_node.PrivateKey().Raw().ToHex() + + sol_ctx = Bip44.FromSeed(seed_bytes, Bip44Coins.SOLANA) + sol_node = ( + sol_ctx.Purpose() + .Coin() + .Account(0) + .Change(Bip44Changes.CHAIN_EXT) + .AddressIndex(0) + ) + sol_addr = sol_node.PublicKey().ToAddress() + sol_pk = sol_node.PrivateKey().Raw().ToHex() + + return [ + DerivedSecretKey( + chain='ETH', + address=eth_addr, + derivation_path=DERIVATION_PATHS['ETH'], + private_key=eth_pk, + ), + DerivedSecretKey( + chain='BSC', + address=eth_addr, + derivation_path=DERIVATION_PATHS['BSC'], + private_key=eth_pk, + ), + DerivedSecretKey( + chain='BTC', + address=btc_addr, + derivation_path=DERIVATION_PATHS['BTC'], + private_key=btc_pk, + ), + DerivedSecretKey( + chain='TRX', + address=trx_addr, + derivation_path=DERIVATION_PATHS['TRX'], + private_key=trx_pk, + ), + DerivedSecretKey( + chain='SOL', + address=sol_addr, + derivation_path=DERIVATION_PATHS['SOL'], + private_key=sol_pk, + ), + ] + + def encrypt_mnemonic(plaintext: str) -> str: if _master_key is None: raise CryptoNotReadyError('Crypto service not ready') diff --git a/src/presentation/dependencies/commands.py b/src/presentation/dependencies/commands.py index f5fec47..b0b0ff5 100644 --- a/src/presentation/dependencies/commands.py +++ b/src/presentation/dependencies/commands.py @@ -10,7 +10,10 @@ from src.application.commands import ( GetAdminMeCommand, CreateOrganizationCommand, CreateOrganizationWalletsCommand, + GetOrganizationMnemonicCommand, + GetOrganizationSecretKeysCommand, GetOrganizationCommand, + ListOrganizationWalletsCommand, GetPurchaseRequestCommand, ListOrganizationsCommand, GetOrganizationDocumentCommand, @@ -91,6 +94,27 @@ def get_create_organization_wallets_command( return CreateOrganizationWalletsCommand(uow, logger) +def get_list_organization_wallets_command( + uow: IUnitOfWork = Depends(get_unit_of_work), + logger: ILogger = Depends(get_logger), +) -> ListOrganizationWalletsCommand: + return ListOrganizationWalletsCommand(uow, logger) + + +def get_get_organization_mnemonic_command( + uow: IUnitOfWork = Depends(get_unit_of_work), + logger: ILogger = Depends(get_logger), +) -> GetOrganizationMnemonicCommand: + return GetOrganizationMnemonicCommand(uow, logger) + + +def get_get_organization_secret_keys_command( + uow: IUnitOfWork = Depends(get_unit_of_work), + logger: ILogger = Depends(get_logger), +) -> GetOrganizationSecretKeysCommand: + return GetOrganizationSecretKeysCommand(uow, logger) + + def get_upload_organization_document_command( uow: IUnitOfWork = Depends(get_unit_of_work), logger: ILogger = Depends(get_logger), diff --git a/src/presentation/routing/organizations.py b/src/presentation/routing/organizations.py index 58d63f6..4cd2549 100644 --- a/src/presentation/routing/organizations.py +++ b/src/presentation/routing/organizations.py @@ -4,6 +4,9 @@ from src.application.commands import ( CreateOrganizationCommand, CreateOrganizationWalletsCommand, GetOrganizationCommand, + GetOrganizationMnemonicCommand, + GetOrganizationSecretKeysCommand, + ListOrganizationWalletsCommand, ListOrganizationsCommand, UpdateOrganizationCommand, ) @@ -13,14 +16,26 @@ from src.presentation.dependencies.commands import ( get_create_organization_command, get_create_organization_wallets_command, get_get_organization_command, + get_get_organization_mnemonic_command, + get_get_organization_secret_keys_command, + get_list_organization_wallets_command, get_list_organizations_command, get_update_organization_command, ) -from src.presentation.schemas.mappers import organization_to_response, wallet_to_response +from src.presentation.schemas.mappers import ( + create_wallets_to_response, + mnemonic_to_response, + organization_to_response, + secret_key_to_response, + wallet_to_response, +) from src.presentation.schemas.organization import ( CreateOrganizationRequest, + CreateWalletsResponse, + MnemonicResponse, OrganizationListResponse, OrganizationResponse, + SecretKeyResponse, UpdateOrganizationRequest, WalletResponse, ) @@ -88,9 +103,47 @@ async def update_organization( return organization_to_response(org) +@organizations_router.get('/{organization_id}/wallets', response_model=list[WalletResponse]) +async def list_organization_wallets( + organization_id: str, + auth: AdminAuthContext = Depends(require_admin_access), + command: ListOrganizationWalletsCommand = Depends(get_list_organization_wallets_command), +): + wallets = await command(organization_id) + return [wallet_to_response(w) for w in wallets] + + +@organizations_router.get('/{organization_id}/wallets/mnemonic', response_model=MnemonicResponse) +async def get_organization_mnemonic( + organization_id: str, + auth: AdminAuthContext = Depends(require_admin_role('compliance', 'superadmin')), + command: GetOrganizationMnemonicCommand = Depends(get_get_organization_mnemonic_command), +): + mnemonic = await command(organization_id) + return mnemonic_to_response(mnemonic) + + +@organizations_router.get('/{organization_id}/wallets/secret-keys', response_model=list[SecretKeyResponse]) +async def get_organization_secret_keys( + organization_id: str, + auth: AdminAuthContext = Depends(require_admin_role('compliance', 'superadmin')), + command: GetOrganizationSecretKeysCommand = Depends(get_get_organization_secret_keys_command), +): + keys = await command(organization_id) + return [ + secret_key_to_response( + chain=k.chain, + address=k.address, + derivation_path=k.derivation_path, + private_key=k.private_key, + ) + for k in keys + ] + + @organizations_router.post( '/{organization_id}/wallets/create', - response_model=list[WalletResponse], + response_model=CreateWalletsResponse, status_code=status.HTTP_201_CREATED, ) async def create_organization_wallets( @@ -98,5 +151,5 @@ async def create_organization_wallets( auth: AdminAuthContext = Depends(require_admin_role('compliance', 'superadmin')), command: CreateOrganizationWalletsCommand = Depends(get_create_organization_wallets_command), ): - wallets = await command(organization_id=organization_id) - return [wallet_to_response(w) for w in wallets] + result = await command(organization_id=organization_id) + return create_wallets_to_response(wallets=result.wallets, mnemonic=result.mnemonic) diff --git a/src/presentation/schemas/mappers.py b/src/presentation/schemas/mappers.py index 4adea44..0f7a83e 100644 --- a/src/presentation/schemas/mappers.py +++ b/src/presentation/schemas/mappers.py @@ -7,9 +7,12 @@ from src.application.domain.entities.organization import ( PurchaseRequestEntity, ) from src.presentation.schemas.organization import ( + CreateWalletsResponse, DocumentResponse, + MnemonicResponse, OrganizationResponse, PurchaseRequestResponse, + SecretKeyResponse, WalletResponse, ) @@ -48,6 +51,32 @@ def wallet_to_response(entity: OrganizationWalletEntity) -> WalletResponse: ) +def create_wallets_to_response(*, wallets: list[OrganizationWalletEntity], mnemonic: str) -> CreateWalletsResponse: + return CreateWalletsResponse( + wallets=[wallet_to_response(w) for w in wallets], + mnemonic=mnemonic, + ) + + +def mnemonic_to_response(mnemonic: str) -> MnemonicResponse: + return MnemonicResponse(mnemonic=mnemonic) + + +def secret_key_to_response( + *, + chain: str, + address: str, + derivation_path: str, + private_key: str, +) -> SecretKeyResponse: + return SecretKeyResponse( + chain=chain, + address=address, + derivation_path=derivation_path, + private_key=private_key, + ) + + def document_to_response( entity: OrganizationDocumentEntity, *, diff --git a/src/presentation/schemas/organization.py b/src/presentation/schemas/organization.py index ee5e0a9..26ab3e2 100644 --- a/src/presentation/schemas/organization.py +++ b/src/presentation/schemas/organization.py @@ -70,6 +70,22 @@ class WalletResponse(BaseModel): created_at: str | None +class CreateWalletsResponse(BaseModel): + wallets: list[WalletResponse] + mnemonic: str + + +class MnemonicResponse(BaseModel): + mnemonic: str + + +class SecretKeyResponse(BaseModel): + chain: str + address: str + derivation_path: str + private_key: str + + class DocumentResponse(BaseModel): id: str organization_id: str