const crypto = require('crypto');

console.log('Custom Diffie-Hellman Key Exchange with Parameter Validation\n');

// Function to validate DH parameters
function validateDHParams(prime, generator) {
  const errors = [];
  
  // Convert to BigInt for mathematical operations
  const p = BigInt('0x' + prime.toString('hex'));
  const g = BigInt('0x' + generator.toString('hex'));
  
  // Check if p is a safe prime (p = 2q + 1 where q is also prime)
  const q = (p - 1n) / 2n;
  
  // Simple primality test (not suitable for production)
  const isPrime = (n) => {
    if (n <= 1n) return false;
    if (n <= 3n) return true;
    if (n % 2n === 0n || n % 3n === 0n) return false;
    
    let i = 5n;
    while (i * i <= n) {
      if (n % i === 0n || n % (i + 2n) === 0n) return false;
      i += 6n;
    }
    return true;
  };
  
  if (!isPrime(p)) errors.push('p is not prime');
  if (!isPrime(q)) errors.push('(p-1)/2 is not prime (p is not a safe prime)');
  if (g <= 1n || g >= p - 1n) errors.push('g must be in the range ]1, p-1[');
  
  return errors.length === 0 ? true : errors;
}

// Create a custom DH group
function createCustomDH(prime, generator) {
  try {
    // Convert to buffers if they're not already
    const primeBuf = Buffer.isBuffer(prime) ? prime : Buffer.from(prime, 'hex');
    const genBuf = Buffer.isBuffer(generator) ? generator : Buffer.from(generator, 'hex');
    
    // Create the DH instance
    const dh = crypto.createDiffieHellman(primeBuf, genBuf);
    
    // Validate parameters
    const validation = validateDHParams(primeBuf, genBuf);
    if (validation !== true) {
      throw new Error(`DiffieHellman parameter validation failed: ${validation.join(', ')}`);
    }
    
    return dh;
  } catch (error) {
    console.error('Error creating custom DH group:', error.message);
    throw error;
  }
}

// Example usage with a small prime (for demonstration only)
// In production, use much larger primes (at least 2048 bits)
const prime = 'f5e108';  // Small prime for demonstration
const generator = '02';  // Common generator value

try {
  console.log('Creating custom DH group...');
  const dh = createCustomDH(prime, generator);
  
  // Generate keys
  dh.generateKeys();
  
  console.log('Custom DH group created successfully');
  console.log('Prime:', dh.getPrime('hex'));
  console.log('Generator:', dh.getGenerator('hex'));
  console.log('Public key:', dh.getPublicKey('hex'));
  console.log('Private key:', dh.getPrivateKey('hex'));
} catch (error) {
  console.error('Failed to create custom DH group:', error.message);
}