const pkcs11 = require("pkcs11js");
const ffi = require("ffi");
const ref = require("ref");
const ArrayType = require('ref-array')
const StructType = require("ref-struct");

const KEYPAIR_ID = Buffer.from("GOST keypair");

// PKCS11 types
const CK_BYTE = ref.types.byte;
const CK_BYTE_PTR = ArrayType(CK_BYTE);
const CK_ULONG = ref.types.long;
const CK_EC_KDF_TYPE = CK_ULONG;

// typedef struct CK_GOSTR3410_DERIVE_PARAMS {
// 	CK_EC_KDF_TYPE kdf;
// 	CK_BYTE_PTR pPublicData;
// 	CK_ULONG ulPublicDataLen;
// 	CK_BYTE_PTR pUKM;
// 	CK_ULONG ulUKMLen;
// } CK_GOSTR3410_DERIVE_PARAMS;

const CK_GOSTR3410_DERIVE_PARAMS = StructType({
  kdf: CK_EC_KDF_TYPE,
  pPublicData: CK_BYTE_PTR,
  ulPublicDataLen: CK_ULONG,
  pUKM: CK_BYTE_PTR,
  ulUKMLen: CK_ULONG,
})

function GenerateGOSTKeyPair(token, session) {
  console.log("Generating GOST key pair");

  const GOST34_10_2001PublicKey = [
    { type: pkcs11.CKA_CLASS, value: pkcs11.CKO_PUBLIC_KEY },
    { type: pkcs11.CKA_LABEL, value: "GOST Public Key" },
    { type: pkcs11.CKA_ID, value: KEYPAIR_ID },
    { type: pkcs11.CKA_KEY_TYPE, value: pkcs11.CKK_GOSTR3410 },
    { type: pkcs11.CKA_TOKEN, value: true },
    { type: pkcs11.CKA_PRIVATE, value: false },
    { type: pkcs11.CKA_DERIVE, value: true },
    { type: pkcs11.CKA_GOSTR3410_PARAMS, value: Buffer.from([0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01]) },
    { type: pkcs11.CKA_GOSTR3411_PARAMS, value: Buffer.from([0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x1e, 0x01]) },
  ];

  const GOST34_10_2001PrivateKey = [
    { type: pkcs11.CKA_CLASS, value: pkcs11.CKO_PRIVATE_KEY },
    { type: pkcs11.CKA_ID, value: KEYPAIR_ID },
    { type: pkcs11.CKA_KEY_TYPE, value: pkcs11.CKK_GOSTR3410 },
    { type: pkcs11.CKA_TOKEN, value: true },
    { type: pkcs11.CKA_PRIVATE, value: true },
    { type: pkcs11.CKA_DERIVE, value: true },
    { type: pkcs11.CKA_GOSTR3410_PARAMS, value: Buffer.from([0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01]) },
  ];

  const keys = token.C_GenerateKeyPair(session,
    { mechanism: pkcs11.CKM_GOSTR3410_KEY_PAIR_GEN, parameter: null },
    GOST34_10_2001PublicKey,
    GOST34_10_2001PrivateKey);

  return keys;
}

function FindKey(token, session, keyClass, id) {
  token.C_FindObjectsInit(session, [
    { type: pkcs11.CKA_CLASS, value: keyClass },
    { type: pkcs11.CKA_ID, value: id },
  ]);

  const obj = token.C_FindObjects(session);
  token.C_FindObjectsFinal(session);

  return obj || null;
}

async function main() {
  const token = new pkcs11.PKCS11();
  token.load("/usr/local/lib/librtpkcs11ecp.dylib");

  token.C_Initialize();

  try {
    console.log(token.C_GetInfo());

    const slots = token.C_GetSlotList(true);
    const slot = slots[0];
    const session = token.C_OpenSession(slot, pkcs11.CKF_SERIAL_SESSION | pkcs11.CKF_RW_SESSION);
    token.C_Login(session, pkcs11.CKU_USER, "12345678");

    // GenerateGOSTKeyPair(token, session);
    const publicKey = FindKey(token, session, pkcs11.CKO_PUBLIC_KEY, KEYPAIR_ID);
    if (!publicKey) throw new Error("Cannot get public key");
    const privateKey = FindKey(token, session, pkcs11.CKO_PRIVATE_KEY, KEYPAIR_ID);
    if (!privateKey) throw new Error("Cannot get private key");

    const attrGOST28147DerivedKey = [
      { type: pkcs11.CKA_CLASS, value: pkcs11.CKO_SECRET_KEY },
      { type: pkcs11.CKA_LABEL, value: "Derived GOST key" },
      { type: pkcs11.CKA_KEY_TYPE, value: pkcs11.CKK_GOST28147 },
      { type: pkcs11.CKA_TOKEN, value: false },
      { type: pkcs11.CKA_PRIVATE, value: false },
      { type: pkcs11.CKA_EXTRACTABLE, value: false },
      { type: pkcs11.CKA_SENSITIVE, value: false },
    ];
    const clientPublicKey = Buffer.from("FF8DAB7F1C0B74A5AD7F0B5F8D5B3C44583798C92586407EEC6EAF00CB4465A5229A53563297358099CA1E17213A960E21FBC60F255B5D994EC45C42087D0604", "hex");
    const UKM = Buffer.from("A93C164618F031F3", "hex");
    const parameters = new CK_GOSTR3410_DERIVE_PARAMS({
      kdf: pkcs11.CKD_CPDIVERSIFY_KDF,
      pPublicData: clientPublicKey,
      ulPublicDataLen: clientPublicKey.length,
      pUKM: UKM,
      ulUKMLen: UKM.length,
    });

    const secKey = token.C_DeriveKey(session,
      { mechanism: pkcs11.CKM_GOSTR3410_DERIVE, parameter: parameters.ref() },
      privateKey,
      attrGOST28147DerivedKey);

    const secKeyValues = token.C_GetAttributeValue(session, secKey, [
      { type: pkcs11.CKA_CLASS, value: null },
      { type: pkcs11.CKA_VALUE, value: null },
    ]);
    console.log("Sec key:", secKeyValues);

  }
  catch (e) {
    console.error(e)
  }
  finally {
    token.C_Finalize();
  }

}

main().catch(err => console.error(err));