smartmta/ts/mail/routing/classes.dkim.manager.ts

import { logger } from '../../logger.js';
import { DKIMCreator } from '../security/classes.dkimcreator.js';
import { DomainRegistry } from './classes.domain.registry.js';
import { RustSecurityBridge } from '../../security/classes.rustsecuritybridge.js';
import { Email } from '../core/classes.email.js';

/** External DcRouter interface shape used by DkimManager */
interface DcRouter {
  storageManager: any;
  dnsServer?: any;
}

/**
 * Manages DKIM key setup, rotation, and signing for all configured domains
 */
export class DkimManager {
  private dkimKeys: Map<string, string> = new Map();

  constructor(
    private dkimCreator: DKIMCreator,
    private domainRegistry: DomainRegistry,
    private dcRouter: DcRouter,
    private rustBridge: RustSecurityBridge,
  ) {}

  async setupDkimForDomains(): Promise<void> {
    const domainConfigs = this.domainRegistry.getAllConfigs();

    if (domainConfigs.length === 0) {
      logger.log('warn', 'No domains configured for DKIM');
      return;
    }

    for (const domainConfig of domainConfigs) {
      const domain = domainConfig.domain;
      const selector = domainConfig.dkim?.selector || 'default';

      try {
        let keyPair: { privateKey: string; publicKey: string };

        try {
          keyPair = await this.dkimCreator.readDKIMKeys(domain);
          logger.log('info', `Using existing DKIM keys for domain: ${domain}`);
        } catch (error) {
          keyPair = await this.dkimCreator.createDKIMKeys();
          await this.dkimCreator.createAndStoreDKIMKeys(domain);
          logger.log('info', `Generated new DKIM keys for domain: ${domain}`);
        }

        this.dkimKeys.set(domain, keyPair.privateKey);
        logger.log('info', `DKIM keys loaded for domain: ${domain} with selector: ${selector}`);
      } catch (error) {
        logger.log('error', `Failed to set up DKIM for domain ${domain}: ${error.message}`);
      }
    }
  }

  async checkAndRotateDkimKeys(): Promise<void> {
    const domainConfigs = this.domainRegistry.getAllConfigs();

    for (const domainConfig of domainConfigs) {
      const domain = domainConfig.domain;
      const selector = domainConfig.dkim?.selector || 'default';
      const rotateKeys = domainConfig.dkim?.rotateKeys || false;
      const rotationInterval = domainConfig.dkim?.rotationInterval || 90;
      const keySize = domainConfig.dkim?.keySize || 2048;

      if (!rotateKeys) {
        logger.log('debug', `DKIM key rotation disabled for ${domain}`);
        continue;
      }

      try {
        const needsRotation = await this.dkimCreator.needsRotation(domain, selector, rotationInterval);

        if (needsRotation) {
          logger.log('info', `DKIM keys need rotation for ${domain} (selector: ${selector})`);

          const newSelector = await this.dkimCreator.rotateDkimKeys(domain, selector, keySize);

          domainConfig.dkim = {
            ...domainConfig.dkim,
            selector: newSelector
          };

          if (domainConfig.dnsMode === 'internal-dns' && this.dcRouter.dnsServer) {
            const keyPair = await this.dkimCreator.readDKIMKeysForSelector(domain, newSelector);
            const publicKeyBase64 = keyPair.publicKey
              .replace(/-----BEGIN PUBLIC KEY-----/g, '')
              .replace(/-----END PUBLIC KEY-----/g, '')
              .replace(/\s/g, '');

            const ttl = domainConfig.dns?.internal?.ttl || 3600;

            this.dcRouter.dnsServer.registerHandler(
              `${newSelector}._domainkey.${domain}`,
              ['TXT'],
              () => ({
                name: `${newSelector}._domainkey.${domain}`,
                type: 'TXT',
                class: 'IN',
                ttl: ttl,
                data: `v=DKIM1; k=rsa; p=${publicKeyBase64}`
              })
            );

            logger.log('info', `DKIM DNS handler registered for new selector: ${newSelector}._domainkey.${domain}`);

            await this.dcRouter.storageManager.set(
              `/email/dkim/${domain}/public.key`,
              keyPair.publicKey
            );
          }

          this.dkimCreator.cleanupOldKeys(domain, 30).catch(error => {
            logger.log('warn', `Failed to cleanup old DKIM keys for ${domain}: ${error.message}`);
          });

        } else {
          logger.log('debug', `DKIM keys for ${domain} are up to date`);
        }
      } catch (error) {
        logger.log('error', `Failed to check/rotate DKIM keys for ${domain}: ${error.message}`);
      }
    }
  }

  async handleDkimSigning(email: Email, domain: string, selector: string): Promise<void> {
    try {
      await this.dkimCreator.handleDKIMKeysForDomain(domain);
      const { privateKey } = await this.dkimCreator.readDKIMKeys(domain);
      const rawEmail = email.toRFC822String();

      // Detect key type from PEM header
      const keyType = privateKey.includes('ED25519') ? 'ed25519' : 'rsa';

      const signResult = await this.rustBridge.signDkim({
        rawMessage: rawEmail,
        domain,
        selector,
        privateKey,
        keyType,
      });

      if (signResult.header) {
        email.addHeader('DKIM-Signature', signResult.header);
        logger.log('info', `Successfully added DKIM signature for ${domain}`);
      }
    } catch (error) {
      logger.log('error', `Failed to sign email with DKIM: ${error.message}`);
    }
  }

  getDkimKey(domain: string): string | undefined {
    return this.dkimKeys.get(domain);
  }
}