#!/usr/bin/env node

/**
 * Rotate an account's signing key onto a new curve.
 *
 * Two writes, in this order. The PLC directory learns the new key first,
 * signed by the rotation key the genesis operation named, which is the key
 * being replaced. The PDS adopts it second, through /init. Doing it the other
 * way would leave the PDS signing with a key the directory does not publish,
 * so nothing could verify its commits until the directory caught up.
 *
 * Usage:
 *   PDS_PASSWORD=... node scripts/rotate-key.js \
 *     --credentials ./credentials-alice.example.com.json \
 *     [--recovery ./recovery-alice.example.com.json] [--curve secp256k1]
 *
 * The new key is written to disk before the directory hears about it. An
 * account whose genesis named only one rotation key has no way back from a
 * key lost between those two writes, so the order is the safety.
 *
 * The credentials file is rewritten, and the old key kept beside it.
 */

import { copyFileSync, readFileSync, rmSync, writeFileSync } from 'node:fs';
import { bytesToHex, importPrivateKey } from '@pdsjs/core/crypto';
import {
  computePlcOperationCid,
  publicKeyToDidKey,
  signPlcOperationObject,
} from '@pdsjs/core/plc';
import { k256Signer } from '@pdsjs/node';

function parseArgs() {
  const args = process.argv.slice(2);
  const opts = {
    credentials: null,
    recovery: null,
    curve: 'secp256k1',
    plcUrl: 'https://plc.directory',
  };
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--credentials' && args[i + 1])
      opts.credentials = args[++i];
    else if (args[i] === '--recovery' && args[i + 1]) opts.recovery = args[++i];
    else if (args[i] === '--curve' && args[i + 1]) opts.curve = args[++i];
    else if (args[i] === '--plc-url' && args[i + 1]) opts.plcUrl = args[++i];
  }
  if (!opts.credentials) {
    console.error(
      'Usage: node scripts/rotate-key.js --credentials <file> [--recovery <file>] [--curve secp256k1|p256]',
    );
    console.error('');
    console.error(
      "  --recovery  The account's recovery key. It outranks the signing key,",
    );
    console.error(
      '              so it can rotate past an operation the signing key cannot.',
    );
    console.error('');
    console.error('Reads PDS_PASSWORD from the environment to authenticate.');
    process.exit(1);
  }
  if (opts.curve !== 'secp256k1' && opts.curve !== 'p256') {
    console.error('--curve must be secp256k1 or p256');
    process.exit(1);
  }
  return opts;
}

/** @param {string} hex */
function hexToBytes(hex) {
  const bytes = new Uint8Array(hex.length / 2);
  for (let i = 0; i < hex.length; i += 2) {
    bytes[i / 2] = Number.parseInt(hex.slice(i, i + 2), 16);
  }
  return bytes;
}

/**
 * The signer for a stored key, on whichever curve it was made.
 * @param {Uint8Array} privateKey
 * @param {'secp256k1'|'p256'} curve
 */
async function signerFor(privateKey, curve) {
  if (curve !== 'secp256k1') return importPrivateKey(privateKey);
  return {
    curve,
    sign: (/** @type {Uint8Array} */ data) => k256Signer.sign(privateKey, data),
    publicKey: async () => k256Signer.getPublicKey(privateKey, true),
  };
}

/** @param {'secp256k1'|'p256'} curve */
async function generateAccountKey(curve) {
  if (curve === 'secp256k1') {
    const privateKey = k256Signer.generatePrivateKey();
    return { privateKey, publicKey: k256Signer.getPublicKey(privateKey, true) };
  }
  const { generateKeyPair } = await import('@pdsjs/core/crypto');
  const keyPair = await generateKeyPair();
  return { privateKey: keyPair.privateKey, publicKey: keyPair.publicKey };
}

/**
 * The operation the new one builds on: the most recent that names the key we
 * hold as a rotation key.
 *
 * Normally that is the head. When it is not, an earlier rotation published a
 * key that never reached the PDS, and building on the last operation we can
 * still sign nullifies it. PLC accepts that for 72 hours, because the key
 * signing it has equal authority to the one that signed what it replaces.
 * @param {string} plcUrl
 * @param {string} did
 * @param {string} ourDidKey
 */
async function fetchPlcState(plcUrl, did, ourDidKey) {
  const res = await fetch(`${plcUrl}/${encodeURIComponent(did)}/log/audit`);
  if (!res.ok) {
    throw new Error(`PLC audit log unavailable: ${res.status}`);
  }
  const log = await res.json();
  const live = log.filter(
    (/** @type {{nullified: boolean}} */ e) => !e.nullified,
  );
  if (live.length === 0) throw new Error('PLC audit log is empty');

  for (let i = live.length - 1; i >= 0; i--) {
    if (live[i].operation.rotationKeys.includes(ourDidKey)) {
      return {
        operation: live[i].operation,
        cid: live[i].cid,
        nullifies: live.length - 1 - i,
      };
    }
  }
  throw new Error(
    `no operation in the PLC log names ${ourDidKey} as a rotation key, so this key cannot rotate the account`,
  );
}

async function main() {
  const opts = parseArgs();
  const creds = JSON.parse(readFileSync(opts.credentials, 'utf8'));
  const oldCurve = creds.curve === 'secp256k1' ? 'secp256k1' : 'p256';

  console.log('Signing key rotation');
  console.log('====================');
  console.log(`DID:   ${creds.did}`);
  console.log(`PDS:   ${creds.pdsUrl}`);
  console.log(`Curve: ${oldCurve} -> ${opts.curve}`);
  console.log('');

  if (oldCurve === opts.curve) {
    console.log(`Already on ${opts.curve}. Nothing to do.`);
    return;
  }

  const password = process.env.PDS_PASSWORD;
  if (!password) {
    console.error('PDS_PASSWORD must be set to authenticate the /init write.');
    process.exit(1);
  }

  console.log(`Generating ${opts.curve} keypair...`);
  const next = await generateAccountKey(opts.curve);
  const nextDidKey = publicKeyToDidKey(next.publicKey, opts.curve);
  console.log(`  did:key: ${nextDidKey}`);

  // The key reaches disk before the directory hears about it. Publishing
  // first and saving after loses the account's rotation key to any failure
  // in between, and the key that could recover it is the one being replaced.
  const pendingFile = `${opts.credentials}.pending`;
  writeFileSync(
    pendingFile,
    `${JSON.stringify(
      {
        did: creds.did,
        curve: opts.curve,
        privateKeyHex: bytesToHex(next.privateKey),
        didKey: nextDidKey,
        note: 'Generated for a rotation that had not finished. Safe to delete once the credentials file names this key.',
      },
      null,
      2,
    )}\n`,
    { mode: 0o600 },
  );
  console.log(`  held at: ${pendingFile}`);
  console.log('');

  // The recovery key outranks the signing key, so it can build on an
  // operation the signing key can no longer sign for.
  const recovery = opts.recovery
    ? JSON.parse(readFileSync(opts.recovery, 'utf8'))
    : null;
  const signerCreds = recovery ?? creds;
  const signerCurve = signerCreds.curve === 'secp256k1' ? 'secp256k1' : 'p256';
  if (recovery) console.log(`Signing with the recovery key ${recovery.didKey}`);

  console.log('Reading the current PLC state...');
  const state = await fetchPlcState(opts.plcUrl, creds.did, signerCreds.didKey);
  console.log(`  prev: ${state.cid}`);
  if (state.nullifies > 0) {
    console.log(
      `  nullifies ${state.nullifies} later operation(s) this key cannot sign for`,
    );
  }
  console.log('');

  // The new key takes both roles the old one held: it verifies the account's
  // commits, and it signs the next rotation.
  const operation = {
    type: 'plc_operation',
    // The recovery key stays, and stays first. A rotation that dropped it
    // would leave the new signing key as the only thing that can rotate
    // again, which is the state this script exists to get out of.
    rotationKeys: recovery
      ? [recovery.didKey, nextDidKey]
      : state.operation.rotationKeys
          .filter((/** @type {string} */ key) => key !== creds.didKey)
          .concat(nextDidKey),
    // Every other verification method, the KEM key among them, stays.
    verificationMethods: {
      ...state.operation.verificationMethods,
      atproto: nextDidKey,
    },
    alsoKnownAs: state.operation.alsoKnownAs,
    services: state.operation.services,
    prev: state.cid,
  };

  console.log('Signing the operation with the current rotation key...');
  const signer = await signerFor(
    hexToBytes(signerCreds.privateKeyHex),
    signerCurve,
  );
  const signed = await signPlcOperationObject(operation, signer);
  console.log(`  cid: ${await computePlcOperationCid(signed)}`);
  console.log('');

  console.log(`Submitting to ${opts.plcUrl}...`);
  const plcRes = await fetch(
    `${opts.plcUrl}/${encodeURIComponent(creds.did)}`,
    {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(signed),
    },
  );
  if (!plcRes.ok) {
    throw new Error(
      `PLC rejected the operation: ${plcRes.status} ${await plcRes.text()}`,
    );
  }
  console.log('  Accepted.');
  console.log('');

  console.log(`Handing the key to ${creds.pdsUrl}...`);
  // An initialized server takes the account's own session here. The password
  // in the body only gates the first bootstrap, when no account exists to
  // authenticate as.
  const sessionRes = await fetch(
    `${creds.pdsUrl}/xrpc/com.atproto.server.createSession`,
    {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ identifier: creds.handle, password }),
    },
  );
  if (!sessionRes.ok) {
    throw new Error(
      `could not sign in to the PDS: ${sessionRes.status} ${await sessionRes.text()}`,
    );
  }
  const { accessJwt } = await sessionRes.json();

  const initRes = await fetch(`${creds.pdsUrl}/init`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${accessJwt}`,
    },
    body: JSON.stringify({
      did: creds.did,
      privateKey: bytesToHex(next.privateKey),
      handle: creds.handle,
      curve: opts.curve,
    }),
  });
  if (!initRes.ok) {
    throw new Error(
      `PDS refused the key: ${initRes.status} ${await initRes.text()}`,
    );
  }
  console.log('  Adopted.');
  console.log('');

  copyFileSync(opts.credentials, `${opts.credentials}.${oldCurve}.bak`);
  writeFileSync(
    opts.credentials,
    `${JSON.stringify(
      {
        ...creds,
        curve: opts.curve,
        privateKeyHex: bytesToHex(next.privateKey),
        didKey: nextDidKey,
        rotatedAt: new Date().toISOString(),
      },
      null,
      2,
    )}\n`,
  );

  rmSync(pendingFile, { force: true });

  console.log('Rotation complete.');
  console.log(`  Credentials rewritten: ${opts.credentials}`);
  console.log(`  Previous key kept at:  ${opts.credentials}.${oldCurve}.bak`);
}

main().catch((err) => {
  console.error(`fatal: ${err instanceof Error ? err.message : err}`);
  process.exit(1);
});