dcrouter/ts/mail/security/classes.dkimcreator.ts

import * as plugins from '../../plugins.js';
import * as paths from '../../paths.js';

import { Email } from '../core/classes.email.js';
// MtaService reference removed

const readFile = plugins.util.promisify(plugins.fs.readFile);
const writeFile = plugins.util.promisify(plugins.fs.writeFile);
const generateKeyPair = plugins.util.promisify(plugins.crypto.generateKeyPair);

export interface IKeyPaths {
  privateKeyPath: string;
  publicKeyPath: string;
}

export interface IDkimKeyMetadata {
  domain: string;
  selector: string;
  createdAt: number;
  rotatedAt?: number;
  previousSelector?: string;
  keySize: number;
}

export class DKIMCreator {
  private keysDir: string;
  private storageManager?: any; // StorageManager instance

  constructor(keysDir = paths.keysDir, storageManager?: any) {
    this.keysDir = keysDir;
    this.storageManager = storageManager;
  }

  public async getKeyPathsForDomain(domainArg: string): Promise<IKeyPaths> {
    return {
      privateKeyPath: plugins.path.join(this.keysDir, `${domainArg}-private.pem`),
      publicKeyPath: plugins.path.join(this.keysDir, `${domainArg}-public.pem`),
    };
  }

  // Check if a DKIM key is present and creates one and stores it to disk otherwise
  public async handleDKIMKeysForDomain(domainArg: string): Promise<void> {
    try {
      await this.readDKIMKeys(domainArg);
    } catch (error) {
      console.log(`No DKIM keys found for ${domainArg}. Generating...`);
      await this.createAndStoreDKIMKeys(domainArg);
      const dnsValue = await this.getDNSRecordForDomain(domainArg);
      plugins.smartfile.fs.ensureDirSync(paths.dnsRecordsDir);
      plugins.smartfile.memory.toFsSync(JSON.stringify(dnsValue, null, 2), plugins.path.join(paths.dnsRecordsDir, `${domainArg}.dkimrecord.json`));
    }
  }

  public async handleDKIMKeysForEmail(email: Email): Promise<void> {
    const domain = email.from.split('@')[1];
    await this.handleDKIMKeysForDomain(domain);
  }

  // Read DKIM keys - always use storage manager, migrate from filesystem if needed
  public async readDKIMKeys(domainArg: string): Promise<{ privateKey: string; publicKey: string }> {
    // Try to read from storage manager first
    if (this.storageManager) {
      try {
        const [privateKey, publicKey] = await Promise.all([
          this.storageManager.get(`/email/dkim/${domainArg}/private.key`),
          this.storageManager.get(`/email/dkim/${domainArg}/public.key`)
        ]);
        
        if (privateKey && publicKey) {
          return { privateKey, publicKey };
        }
      } catch (error) {
        // Fall through to migration check
      }
      
      // Check if keys exist in filesystem and migrate them to storage manager
      const keyPaths = await this.getKeyPathsForDomain(domainArg);
      try {
        const [privateKeyBuffer, publicKeyBuffer] = await Promise.all([
          readFile(keyPaths.privateKeyPath),
          readFile(keyPaths.publicKeyPath),
        ]);

        // Convert the buffers to strings
        const privateKey = privateKeyBuffer.toString();
        const publicKey = publicKeyBuffer.toString();
        
        // Migrate to storage manager
        console.log(`Migrating DKIM keys for ${domainArg} from filesystem to StorageManager`);
        await Promise.all([
          this.storageManager.set(`/email/dkim/${domainArg}/private.key`, privateKey),
          this.storageManager.set(`/email/dkim/${domainArg}/public.key`, publicKey)
        ]);
        
        return { privateKey, publicKey };
      } catch (error) {
        if (error.code === 'ENOENT') {
          // Keys don't exist anywhere
          throw new Error(`DKIM keys not found for domain ${domainArg}`);
        }
        throw error;
      }
    } else {
      // No storage manager, use filesystem directly
      const keyPaths = await this.getKeyPathsForDomain(domainArg);
      const [privateKeyBuffer, publicKeyBuffer] = await Promise.all([
        readFile(keyPaths.privateKeyPath),
        readFile(keyPaths.publicKeyPath),
      ]);

      const privateKey = privateKeyBuffer.toString();
      const publicKey = publicKeyBuffer.toString();
      
      return { privateKey, publicKey };
    }
  }

  // Create a DKIM key pair - changed to public for API access
  public async createDKIMKeys(): Promise<{ privateKey: string; publicKey: string }> {
    const { privateKey, publicKey } = await generateKeyPair('rsa', {
      modulusLength: 2048,
      publicKeyEncoding: { type: 'spki', format: 'pem' },
      privateKeyEncoding: { type: 'pkcs1', format: 'pem' },
    });

    return { privateKey, publicKey };
  }

  // Store a DKIM key pair - uses storage manager if available, else disk
  public async storeDKIMKeys(
    privateKey: string,
    publicKey: string,
    privateKeyPath: string,
    publicKeyPath: string
  ): Promise<void> {
    // Store in storage manager if available
    if (this.storageManager) {
      // Extract domain from path (e.g., /path/to/keys/example.com-private.pem -> example.com)
      const match = privateKeyPath.match(/\/([^\/]+)-private\.pem$/);
      if (match) {
        const domain = match[1];
        await Promise.all([
          this.storageManager.set(`/email/dkim/${domain}/private.key`, privateKey),
          this.storageManager.set(`/email/dkim/${domain}/public.key`, publicKey)
        ]);
      }
    }
    
    // Also store to filesystem for backward compatibility
    await Promise.all([writeFile(privateKeyPath, privateKey), writeFile(publicKeyPath, publicKey)]);
  }

  // Create a DKIM key pair and store it to disk - changed to public for API access
  public async createAndStoreDKIMKeys(domain: string): Promise<void> {
    const { privateKey, publicKey } = await this.createDKIMKeys();
    const keyPaths = await this.getKeyPathsForDomain(domain);
    await this.storeDKIMKeys(
      privateKey,
      publicKey,
      keyPaths.privateKeyPath,
      keyPaths.publicKeyPath
    );
    console.log(`DKIM keys for ${domain} created and stored.`);
  }

  // Changed to public for API access
  public async getDNSRecordForDomain(domainArg: string): Promise<plugins.tsclass.network.IDnsRecord> {
    await this.handleDKIMKeysForDomain(domainArg);
    const keys = await this.readDKIMKeys(domainArg);

    // Remove the PEM header and footer and newlines
    const pemHeader = '-----BEGIN PUBLIC KEY-----';
    const pemFooter = '-----END PUBLIC KEY-----';
    const keyContents = keys.publicKey
      .replace(pemHeader, '')
      .replace(pemFooter, '')
      .replace(/\n/g, '');

    // Now generate the DKIM DNS TXT record
    const dnsRecordValue = `v=DKIM1; h=sha256; k=rsa; p=${keyContents}`;

    return {
      name: `mta._domainkey.${domainArg}`,
      type: 'TXT',
      dnsSecEnabled: null,
      value: dnsRecordValue,
    };
  }

  /**
   * Get DKIM key metadata for a domain
   */
  private async getKeyMetadata(domain: string, selector: string = 'default'): Promise<IDkimKeyMetadata | null> {
    if (!this.storageManager) {
      return null;
    }
    
    const metadataKey = `/email/dkim/${domain}/${selector}/metadata`;
    const metadataStr = await this.storageManager.get(metadataKey);
    
    if (!metadataStr) {
      return null;
    }
    
    return JSON.parse(metadataStr) as IDkimKeyMetadata;
  }

  /**
   * Save DKIM key metadata
   */
  private async saveKeyMetadata(metadata: IDkimKeyMetadata): Promise<void> {
    if (!this.storageManager) {
      return;
    }
    
    const metadataKey = `/email/dkim/${metadata.domain}/${metadata.selector}/metadata`;
    await this.storageManager.set(metadataKey, JSON.stringify(metadata));
  }

  /**
   * Check if DKIM keys need rotation
   */
  public async needsRotation(domain: string, selector: string = 'default', rotationIntervalDays: number = 90): Promise<boolean> {
    const metadata = await this.getKeyMetadata(domain, selector);
    
    if (!metadata) {
      // No metadata means old keys, should rotate
      return true;
    }
    
    const now = Date.now();
    const keyAgeMs = now - metadata.createdAt;
    const keyAgeDays = keyAgeMs / (1000 * 60 * 60 * 24);
    
    return keyAgeDays >= rotationIntervalDays;
  }

  /**
   * Rotate DKIM keys for a domain
   */
  public async rotateDkimKeys(domain: string, currentSelector: string = 'default', keySize: number = 2048): Promise<string> {
    console.log(`Rotating DKIM keys for ${domain}...`);
    
    // Generate new selector based on date
    const now = new Date();
    const newSelector = `key${now.getFullYear()}${String(now.getMonth() + 1).padStart(2, '0')}`;
    
    // Create new keys with custom key size
    const { privateKey, publicKey } = await generateKeyPair('rsa', {
      modulusLength: keySize,
      publicKeyEncoding: { type: 'spki', format: 'pem' },
      privateKeyEncoding: { type: 'pkcs1', format: 'pem' },
    });
    
    // Store new keys with new selector
    const newKeyPaths = await this.getKeyPathsForSelector(domain, newSelector);
    
    // Store in storage manager if available
    if (this.storageManager) {
      await Promise.all([
        this.storageManager.set(`/email/dkim/${domain}/${newSelector}/private.key`, privateKey),
        this.storageManager.set(`/email/dkim/${domain}/${newSelector}/public.key`, publicKey)
      ]);
    }
    
    // Also store to filesystem
    await this.storeDKIMKeys(
      privateKey,
      publicKey,
      newKeyPaths.privateKeyPath,
      newKeyPaths.publicKeyPath
    );
    
    // Save metadata for new keys
    const metadata: IDkimKeyMetadata = {
      domain,
      selector: newSelector,
      createdAt: Date.now(),
      previousSelector: currentSelector,
      keySize
    };
    await this.saveKeyMetadata(metadata);
    
    // Update metadata for old keys
    const oldMetadata = await this.getKeyMetadata(domain, currentSelector);
    if (oldMetadata) {
      oldMetadata.rotatedAt = Date.now();
      await this.saveKeyMetadata(oldMetadata);
    }
    
    console.log(`DKIM keys rotated for ${domain}. New selector: ${newSelector}`);
    return newSelector;
  }

  /**
   * Get key paths for a specific selector
   */
  public async getKeyPathsForSelector(domain: string, selector: string): Promise<IKeyPaths> {
    return {
      privateKeyPath: plugins.path.join(this.keysDir, `${domain}-${selector}-private.pem`),
      publicKeyPath: plugins.path.join(this.keysDir, `${domain}-${selector}-public.pem`),
    };
  }

  /**
   * Read DKIM keys for a specific selector
   */
  public async readDKIMKeysForSelector(domain: string, selector: string): Promise<{ privateKey: string; publicKey: string }> {
    // Try to read from storage manager first
    if (this.storageManager) {
      try {
        const [privateKey, publicKey] = await Promise.all([
          this.storageManager.get(`/email/dkim/${domain}/${selector}/private.key`),
          this.storageManager.get(`/email/dkim/${domain}/${selector}/public.key`)
        ]);
        
        if (privateKey && publicKey) {
          return { privateKey, publicKey };
        }
      } catch (error) {
        // Fall through to migration check
      }
      
      // Check if keys exist in filesystem and migrate them to storage manager
      const keyPaths = await this.getKeyPathsForSelector(domain, selector);
      try {
        const [privateKeyBuffer, publicKeyBuffer] = await Promise.all([
          readFile(keyPaths.privateKeyPath),
          readFile(keyPaths.publicKeyPath),
        ]);

        const privateKey = privateKeyBuffer.toString();
        const publicKey = publicKeyBuffer.toString();
        
        // Migrate to storage manager
        console.log(`Migrating DKIM keys for ${domain}/${selector} from filesystem to StorageManager`);
        await Promise.all([
          this.storageManager.set(`/email/dkim/${domain}/${selector}/private.key`, privateKey),
          this.storageManager.set(`/email/dkim/${domain}/${selector}/public.key`, publicKey)
        ]);
        
        return { privateKey, publicKey };
      } catch (error) {
        if (error.code === 'ENOENT') {
          throw new Error(`DKIM keys not found for domain ${domain} with selector ${selector}`);
        }
        throw error;
      }
    } else {
      // No storage manager, use filesystem directly
      const keyPaths = await this.getKeyPathsForSelector(domain, selector);
      const [privateKeyBuffer, publicKeyBuffer] = await Promise.all([
        readFile(keyPaths.privateKeyPath),
        readFile(keyPaths.publicKeyPath),
      ]);

      const privateKey = privateKeyBuffer.toString();
      const publicKey = publicKeyBuffer.toString();

      return { privateKey, publicKey };
    }
  }

  /**
   * Get DNS record for a specific selector
   */
  public async getDNSRecordForSelector(domain: string, selector: string): Promise<plugins.tsclass.network.IDnsRecord> {
    const keys = await this.readDKIMKeysForSelector(domain, selector);

    // Remove the PEM header and footer and newlines
    const pemHeader = '-----BEGIN PUBLIC KEY-----';
    const pemFooter = '-----END PUBLIC KEY-----';
    const keyContents = keys.publicKey
      .replace(pemHeader, '')
      .replace(pemFooter, '')
      .replace(/\n/g, '');

    // Generate the DKIM DNS TXT record
    const dnsRecordValue = `v=DKIM1; h=sha256; k=rsa; p=${keyContents}`;

    return {
      name: `${selector}._domainkey.${domain}`,
      type: 'TXT',
      dnsSecEnabled: null,
      value: dnsRecordValue,
    };
  }

  /**
   * Clean up old DKIM keys after grace period
   */
  public async cleanupOldKeys(domain: string, gracePeriodDays: number = 30): Promise<void> {
    if (!this.storageManager) {
      return;
    }
    
    // List all selectors for the domain
    const metadataKeys = await this.storageManager.list(`/email/dkim/${domain}/`);
    
    for (const key of metadataKeys) {
      if (key.endsWith('/metadata')) {
        const metadataStr = await this.storageManager.get(key);
        if (metadataStr) {
          const metadata = JSON.parse(metadataStr) as IDkimKeyMetadata;
          
          // Check if key is rotated and past grace period
          if (metadata.rotatedAt) {
            const gracePeriodMs = gracePeriodDays * 24 * 60 * 60 * 1000;
            const now = Date.now();
            
            if (now - metadata.rotatedAt > gracePeriodMs) {
              console.log(`Cleaning up old DKIM keys for ${domain} selector ${metadata.selector}`);
              
              // Delete key files
              const keyPaths = await this.getKeyPathsForSelector(domain, metadata.selector);
              try {
                await plugins.fs.promises.unlink(keyPaths.privateKeyPath);
                await plugins.fs.promises.unlink(keyPaths.publicKeyPath);
              } catch (error) {
                console.warn(`Failed to delete old key files: ${error.message}`);
              }
              
              // Delete metadata
              await this.storageManager.delete(key);
            }
          }
        }
      }
    }
  }
}