rwadurian/backend/services/identity-service/src/infrastructure/external/blockchain/wallet-generator.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { HDKey } from '@scure/bip32';
import {
  createHash,
  createCipheriv,
  createDecipheriv,
  randomBytes,
  scryptSync,
} from 'crypto';
import { bech32 } from 'bech32';
import { ethers } from 'ethers';
import { ConfigService } from '@nestjs/config';
import { Mnemonic, UserId, ChainType, CHAIN_CONFIG } from '@/domain/value-objects';
import { WalletAddress } from '@/domain/entities/wallet-address.entity';
import { DomainError } from '@/shared/exceptions/domain.exception';

export interface WalletSystemResult {
  mnemonic: Mnemonic;
  wallets: Map<ChainType, WalletAddress>;
}

export interface EncryptedMnemonicData {
  encrypted: string;
  authTag: string;
  iv: string;
}

@Injectable()
export class WalletGeneratorService {
  private readonly logger = new Logger(WalletGeneratorService.name);
  private readonly encryptionSalt: string;

  constructor(private readonly configService: ConfigService) {
    this.encryptionSalt = configService.get(
      'WALLET_ENCRYPTION_SALT',
      'rwa-wallet-salt',
    );
  }

  generateWalletSystem(params: {
    userId: string;
    deviceId: string;
  }): WalletSystemResult {
    const mnemonic = Mnemonic.generate();
    const encryptionKey = this.deriveEncryptionKey(
      params.deviceId,
      params.userId,
    );

    const wallets = new Map<ChainType, WalletAddress>();
    const chains = [ChainType.KAVA, ChainType.DST, ChainType.BSC];
    const userId = UserId.create(params.userId);

    for (const chainType of chains) {
      const wallet = WalletAddress.createFromMnemonic({
        userId,
        chainType,
        mnemonic,
        encryptionKey,
      });

      wallets.set(chainType, wallet);
    }

    this.logger.debug(`Generated wallet system for user: ${params.userId}`);

    return { mnemonic, wallets };
  }

  recoverWalletSystem(params: {
    userId: string;
    mnemonic: Mnemonic;
    deviceId: string;
  }): Map<ChainType, WalletAddress> {
    const encryptionKey = this.deriveEncryptionKey(
      params.deviceId,
      params.userId,
    );

    const wallets = new Map<ChainType, WalletAddress>();
    const chains = [ChainType.KAVA, ChainType.DST, ChainType.BSC];
    const userId = UserId.create(params.userId);

    for (const chainType of chains) {
      const wallet = WalletAddress.createFromMnemonic({
        userId,
        chainType,
        mnemonic: params.mnemonic,
        encryptionKey,
      });

      wallets.set(chainType, wallet);
    }

    this.logger.debug(`Recovered wallet system for user: ${params.userId}`);

    return wallets;
  }

  deriveAddress(chainType: ChainType, mnemonic: Mnemonic): string {
    const seed = Buffer.from(mnemonic.toSeed());
    const config = CHAIN_CONFIG[chainType];

    switch (chainType) {
      case ChainType.KAVA:
      case ChainType.DST:
        return this.deriveCosmosAddress(
          seed,
          config.derivationPath,
          config.prefix,
        );

      case ChainType.BSC:
        return this.deriveEVMAddress(seed, config.derivationPath);

      default:
        throw new DomainError(`不支持的链类型: ${chainType}`);
    }
  }

  verifyMnemonic(
    mnemonic: Mnemonic,
    chainType: ChainType,
    expectedAddress: string,
  ): boolean {
    const derivedAddress = this.deriveAddress(chainType, mnemonic);
    return derivedAddress.toLowerCase() === expectedAddress.toLowerCase();
  }

  private deriveCosmosAddress(
    seed: Buffer,
    path: string,
    prefix: string,
  ): string {
    const hdkey = HDKey.fromMasterSeed(seed);
    const childKey = hdkey.derive(path);

    if (!childKey.publicKey) {
      throw new DomainError('无法派生公钥');
    }

    const pubkey = childKey.publicKey;
    const hash = createHash('sha256').update(pubkey).digest();
    const addressHash = createHash('ripemd160').update(hash).digest();
    const words = bech32.toWords(addressHash);

    return bech32.encode(prefix, words);
  }

  private deriveEVMAddress(seed: Buffer, path: string): string {
    const hdkey = HDKey.fromMasterSeed(seed);
    const childKey = hdkey.derive(path);

    if (!childKey.privateKey) {
      throw new DomainError('无法派生私钥');
    }

    // 将 Uint8Array 转换为十六进制字符串
    const privateKeyHex = '0x' + Buffer.from(childKey.privateKey).toString('hex');
    const wallet = new ethers.Wallet(privateKeyHex);
    return wallet.address;
  }

  encryptMnemonic(mnemonic: string, key: string): EncryptedMnemonicData {
    const derivedKey = scryptSync(key, this.encryptionSalt, 32);
    const iv = randomBytes(16);
    const cipher = createCipheriv('aes-256-gcm', derivedKey, iv);

    let encrypted = cipher.update(mnemonic, 'utf8', 'hex');
    encrypted += cipher.final('hex');
    const authTag = cipher.getAuthTag();

    return {
      encrypted,
      authTag: authTag.toString('hex'),
      iv: iv.toString('hex'),
    };
  }

  decryptMnemonic(
    encryptedData: EncryptedMnemonicData,
    key: string,
  ): string {
    const derivedKey = scryptSync(key, this.encryptionSalt, 32);
    const iv = Buffer.from(encryptedData.iv, 'hex');
    const authTag = Buffer.from(encryptedData.authTag, 'hex');
    const decipher = createDecipheriv('aes-256-gcm', derivedKey, iv);
    decipher.setAuthTag(authTag);

    let decrypted = decipher.update(encryptedData.encrypted, 'hex', 'utf8');
    decrypted += decipher.final('utf8');

    return decrypted;
  }

  deriveEncryptionKey(deviceId: string, userId: string): string {
    const input = `${deviceId}:${userId}`;
    return createHash('sha256').update(input).digest('hex');
  }
}