111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
from os import getenv, urandom
|
|
import os
|
|
import time
|
|
import json
|
|
|
|
from nacl.bindings import crypto_sign_seed_keypair
|
|
from tonsdk.utils import Address
|
|
|
|
from app.core._blockchain.ton.wallet_v3cr3 import WalletV3CR3
|
|
from app.core.logger import make_log
|
|
from sqlalchemy import create_engine, text
|
|
from typing import Optional
|
|
|
|
|
|
def _load_seed_from_env_or_generate() -> bytes:
|
|
seed_hex = os.getenv("TON_INIT_HOT_SEED")
|
|
if seed_hex:
|
|
make_log("HotWallet", "Loaded seed from env")
|
|
return bytes.fromhex(seed_hex)
|
|
make_log("HotWallet", "No seed provided; generating ephemeral seed", level='info')
|
|
return urandom(32)
|
|
|
|
|
|
def _init_seed_via_db() -> bytes:
|
|
"""Store and read hot seed from PostgreSQL service_config (key='private_key').
|
|
Primary node writes it once; workers wait until it appears.
|
|
"""
|
|
from app.core._config import DATABASE_URL
|
|
|
|
engine = create_engine(DATABASE_URL, pool_pre_ping=True)
|
|
role = os.getenv("NODE_ROLE", "worker").lower()
|
|
|
|
def db_ready(conn) -> bool:
|
|
try:
|
|
r = conn.execute(text("SELECT to_regclass('public.service_config')")).scalar()
|
|
return r is not None
|
|
except Exception:
|
|
return False
|
|
|
|
# Wait for table to exist
|
|
start = time.time()
|
|
with engine.connect() as conn:
|
|
while not db_ready(conn):
|
|
time.sleep(0.5)
|
|
if time.time() - start > 120:
|
|
raise TimeoutError("service_config table not available")
|
|
|
|
def read_seed() -> Optional[bytes]:
|
|
row = conn.execute(text("SELECT packed_value FROM service_config WHERE key = :k LIMIT 1"), {"k": "private_key"}).first()
|
|
if not row:
|
|
return None
|
|
packed = row[0] or {}
|
|
try:
|
|
# packed_value is JSON; ensure dict
|
|
if isinstance(packed, str):
|
|
packed = json.loads(packed)
|
|
seed_hex = packed.get("value")
|
|
return bytes.fromhex(seed_hex) if seed_hex else None
|
|
except Exception:
|
|
return None
|
|
|
|
seed = read_seed()
|
|
if seed:
|
|
return seed
|
|
|
|
if role == "primary":
|
|
seed = _load_seed_from_env_or_generate()
|
|
# Try insert; if another primary raced, ignore
|
|
try:
|
|
conn.execute(
|
|
text("INSERT INTO service_config (key, packed_value) VALUES (:k, CAST(:v AS JSON))"),
|
|
{"k": "private_key", "v": json.dumps({"value": seed.hex()})}
|
|
)
|
|
conn.commit()
|
|
make_log("HotWallet", "Seed saved in service_config by primary", level='info')
|
|
return seed
|
|
except Exception:
|
|
conn.rollback()
|
|
# Read again in case of race
|
|
seed2 = read_seed()
|
|
if seed2:
|
|
return seed2
|
|
raise
|
|
else:
|
|
make_log("HotWallet", "Worker waiting for seed in service_config...", level='info')
|
|
while True:
|
|
seed = read_seed()
|
|
if seed:
|
|
return seed
|
|
time.sleep(0.5)
|
|
|
|
|
|
_extra_ton_wallet_options = {}
|
|
if getenv('TON_CUSTOM_WALLET_ADDRESS'):
|
|
_extra_ton_wallet_options['address'] = Address(getenv('TON_CUSTOM_WALLET_ADDRESS'))
|
|
|
|
|
|
def _init_wallet():
|
|
# Primary writes to DB; workers wait and read from DB
|
|
hot_seed_bytes = _init_seed_via_db()
|
|
pub, priv = crypto_sign_seed_keypair(hot_seed_bytes)
|
|
wallet = WalletV3CR3(
|
|
private_key=priv,
|
|
public_key=pub,
|
|
**_extra_ton_wallet_options
|
|
)
|
|
return hot_seed_bytes, pub, priv, wallet
|
|
|
|
|
|
hot_seed, hot_pubkey, hot_privkey, service_wallet = _init_wallet()
|