import assert from 'assert';
import crypto from 'crypto';
import { promisify } from 'util';

// The format of the encrypted data generated by this module is
// version || salt || iv || authTag || ciphertext.
// The version field is included in the encrypted data as well so that some
// properties of the encryption algorithm can be exchanged independently of
// the version (e.g. some other algorithm than PBKDF2 with 100k iterations).
const kEncryptVersionOffset = 0;
const kEncryptVersionLength = 2;
const kSaltOffset = kEncryptVersionOffset + kEncryptVersionLength;
const kSaltLength = 16;
const kIvOffset = kSaltOffset + kSaltLength;
const kIvLength = 16;
const kAuthTagOffset = kIvOffset + kIvLength;
const kAuthTagLength = 16;
const kCiphertextOffset = kAuthTagOffset + kAuthTagLength;

const randomBytes = promisify(crypto.randomBytes);
const pbkdf2 = promisify(crypto.pbkdf2);

/**
 * Encrypter class for encrypting data with a user-specified passphrase.
 * The derived key is cached across encrypt() calls.
 */
export class Encrypter {
  private _passphrase: string;
  private _cachedKey?: { salt: Buffer; key: Buffer };
  private _encryptVersionBuffer = Buffer.from([0x00, 0x01]);

  constructor(passphrase: string) {
    this._passphrase = passphrase;
  }

  private async _getKey(): Promise<{ salt: Buffer; key: Buffer }> {
    if (this._cachedKey) {
      return this._cachedKey;
    }

    // pbkdf2 is the most computationally heavy part of encryption, so caching
    // the key saves a lot of time on repeated encrypt() calls.
    const salt = await randomBytes(kSaltLength);
    const key = await pbkdf2(
      Buffer.from(this._passphrase, 'utf8'),
      salt,
      100_000,
      32,
      'sha512'
    );
    return (this._cachedKey = { salt, key });
  }

  /**
   * Encrypt `plaintext` using the pre-specified passphrase.
   *
   * @param plaintext Any string.
   * @returns A base64-encoded encrypted string.
   */
  async encrypt(plaintext: string): Promise<string> {
    const { salt, key } = await this._getKey();

    const iv = await randomBytes(kIvLength);

    const cipher = crypto.createCipheriv('aes-256-gcm', key, iv, {
      authTagLength: kAuthTagLength,
    });
    const ciphertext = [];
    ciphertext.push(cipher.update(Buffer.from(plaintext, 'utf8')));
    ciphertext.push(cipher.final());
    const authTag = cipher.getAuthTag();

    // Make sure that lengths are what we expect them to be.
    // This should never really fail as the algorithm's used here don't change,
    // but the potential risk of not being able to decrypt data later
    // probably justifies extra checks.
    assert.deepStrictEqual(
      [
        this._encryptVersionBuffer.length,
        salt.length,
        iv.length,
        authTag.length,
      ],
      [kEncryptVersionLength, kSaltLength, kIvLength, kAuthTagLength]
    );

    return Buffer.concat([
      this._encryptVersionBuffer,
      salt,
      iv,
      authTag,
      ...ciphertext,
    ]).toString('base64');
  }
}

/**
 * Encrypter class for encrypting data with a user-specified passphrase.
 */
export class Decrypter {
  private _passphrase: string;
  private _cachedKeys = new Map<string, Buffer>();

  constructor(passphrase: string) {
    this._passphrase = passphrase;
  }

  private async _getKey(salt: Buffer): Promise<Buffer> {
    // The derived key differs for different salts, so we
    // need to potentially cache multiple derived keys here.
    const cacheKey = salt.toString('base64');
    let key = this._cachedKeys.get(cacheKey);
    if (!key) {
      key = await pbkdf2(
        Buffer.from(this._passphrase, 'utf8'),
        salt,
        100_000,
        32,
        'sha512'
      );
      this._cachedKeys.set(cacheKey, key);
    }
    return key;
  }

  /**
   * Decrypt a previously encrypted string using the pre-specified passphrase.
   *
   * @param encryptedBase64 A string returned from @ref Encrypter.encrypt().
   * @returns The original plaintext data.
   */
  async decrypt(encryptedBase64: string): Promise<string> {
    const encrypted = Buffer.from(encryptedBase64, 'base64');
    const version = encrypted
      .subarray(
        kEncryptVersionOffset,
        kEncryptVersionOffset + kEncryptVersionLength
      )
      .readUInt16BE();
    const salt = encrypted.subarray(kSaltOffset, kSaltOffset + kSaltLength);
    const iv = encrypted.subarray(kIvOffset, kIvOffset + kIvLength);
    const authTag = encrypted.subarray(
      kAuthTagOffset,
      kAuthTagOffset + kAuthTagLength
    );
    const ciphertext = encrypted.subarray(kCiphertextOffset);

    if (version !== 0x0001) {
      throw new Error(
        `Unknown encrypted data version: 0x${version.toString(16)}`
      );
    }

    const key = await this._getKey(salt);
    const cipher = crypto.createDecipheriv('aes-256-gcm', key, iv);
    cipher.setAuthTag(authTag);
    const plaintext = [];
    plaintext.push(cipher.update(ciphertext));
    try {
      plaintext.push(cipher.final());
    } catch (err: any) {
      if (err?.message.includes('unable to authenticate data')) {
        throw new Error(
          'Cannot decrypt due to corrupt data or wrong passphrase'
        );
      }
      throw err;
    }

    return Buffer.concat(plaintext).toString('utf8');
  }
}