chore: refactor dkim script for new version

2026-01-09 13:38:03 -05:00 · 2023-11-06 18:52:40 +03:00
parent 43927dfcd9
commit ea0a2aa011
7 changed files with 312 additions and 210 deletions
--- a/packages/helpers/src/dkim/pull-and-format-dkim-key.ts
+++ b/packages/helpers/src/dkim/pull-and-format-dkim-key.ts
@@ -1,70 +0,0 @@
-import dns from "dns";
-import forge from "node-forge";
-import { publicEncrypt } from 'crypto';
-import {
-    toCircomBigIntBytes,
-  } from "../binaryFormat";
-import { pki } from "node-forge";
-
-// Fetch the DKIM public key from DNS and format it for circom
-// Does not output the hash, only outputs it split into 17 parts
-export default async function formatDkimKey(domain: string, selector: string, print: boolean = true) {
-  // Construct the DKIM record name
-  let dkimRecordName = `${selector}._domainkey.${domain}`;
-  if (print) console.log(dkimRecordName);
-  // Lookup the DKIM record in DNS
-  let records;
-  try {
-    records = await dns.promises.resolveTxt(dkimRecordName);
-  } catch (err) {
-    if (print) console.error(err);
-    return;
-  }
-
-  // The DKIM record is a TXT record containing a string
-  // We need to parse this string to get the public key
-  let dkimRecord = records[0].join("");
-  let match = dkimRecord.match(/p=([^;]+)/);
-  if (!match) {
-    console.error("No public key found in DKIM record");
-    return;
-  }
-
-  // The public key is base64 encoded, we need to decode it
-  let pubkey = match[1];
-  let binaryKey = Buffer.from(pubkey, "base64").toString('base64');
-
-  // Get match
-  let matches = binaryKey.match(/.{1,64}/g);
-  if (!matches) {
-    console.error("No matches found");
-    return;
-  }
-  let formattedKey = matches.join("\n");
-  if (print) console.log("Key: ", formattedKey);
-
-  // Convert to PEM format
-  let pemKey = `-----BEGIN PUBLIC KEY-----\n${formattedKey}\n-----END PUBLIC KEY-----`;
-
-  // Parse the RSA public key
-  let publicKey = forge.pki.publicKeyFromPem(pemKey);
-
-  // Get the modulus n only
-  let n = publicKey.n;
-  if (print) console.log("Modulus n:", n.toString(16));
-
-  // Convert binary to BigInt
-  let bigIntKey = BigInt(publicKey.n.toString());
-  if (print) console.log(bigIntKey);
-  if (print) console.log(toCircomBigIntBytes(bigIntKey));
-  return toCircomBigIntBytes(bigIntKey);
-}
-
-// formatDkimKey("gmail.com", "20221208");
-// formatDkimKey("gmail.com", "20230601");
-// formatDkimKey("twitter.com", "dkim-201406");
-// formatDkimKey("ethereum.org", "salesforceeth123321");
-// // let pubkey = "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAq8JxVBMLHZRj1WvIMSHApRY3DraE/EiFiR6IMAlDq9GAnrVy0tDQyBND1G8+1fy5RwssQ9DgfNe7rImwxabWfWxJ1LSmo/DzEdOHOJNQiP/nw7MdmGu+R9hEvBeGRQAmn1jkO46KIw/p2lGvmPSe3+AVD+XyaXZ4vJGTZKFUCnoctAVUyHjSDT7KnEsaiND2rVsDvyisJUAH+EyRfmHSBwfJVHAdJ9oD8cn9NjIun/EHLSIwhCxXmLJlaJeNAFtcGeD2aRGbHaS7M6aTFP+qk4f2ucRx31cyCxbu50CDVfU+d4JkIDNBFDiV+MIpaDFXIf11bGoS08oBBQiyPXgX0wIDAQAB";
-// let bigIntKey = BigInt('0x' + binaryKey.toString('hex'));
-// let binaryKey = Buffer.from(pubkey, 'base64');
-// console.log(toCircomBigIntBytes(bigIntKey));
--- a/packages/helpers/src/dkim/update-dkim-registry.ts
+++ b/packages/helpers/src/dkim/update-dkim-registry.ts
@@ -1,140 +0,0 @@
-import { ethers, AlchemyProvider, InfuraProvider } from 'ethers';
-import formatDkimKey from './pull-and-format-dkim-key';
-import { readFileSync } from 'fs';
-
-require("dotenv").config();
-
-const network = 'goerli'; // or whatever network you're using
-const alchemyApiKey = process.env.ALCHEMY_GOERLI_KEY;
-const infuraApiKey = process.env.INFURA_KEY;
-const localSecretKey = process.env.PRIVATE_KEY || '0';
-const mailserver_address = "0x638E55F942cBD6f8cb715e9C6a9d747c6F852196";
-
-const default_abi = [{
-    "inputs": [
-      {
-        "internalType": "string",
-        "name": "domain",
-        "type": "string"
-      },
-      {
-        "internalType": "uint256",
-        "name": "index",
-        "type": "uint256"
-      },
-      {
-        "internalType": "uint256",
-        "name": "val",
-        "type": "uint256"
-      }
-    ],
-    "name": "editMailserverKey",
-    "outputs": [],
-    "stateMutability": "nonpayable",
-    "type": "function"
-  }];
-
-async function updateMailserverKeys(domain: string, selector: string, contract_address: string, abi: any = default_abi, parallel = true) {
-  // const provider = new AlchemyProvider(network, alchemyApiKey);
-  const provider = new InfuraProvider(network, infuraApiKey);
-  const wallet = new ethers.Wallet(localSecretKey, provider);
-  const contract = new ethers.Contract(contract_address, abi, wallet);
-
-  const publicKeyParts = await formatDkimKey(domain, selector, false);
-
-  if (!publicKeyParts) {
-      console.log('No public key found');
-      return;
-  }
-
-  if (parallel) {
-    let nonce = await provider.getTransactionCount(wallet.address);
-    const txs = publicKeyParts.map(async (part, i) => {
-      const tx = await contract.editMailserverKey(domain, i, part, { nonce: nonce++ });
-      return tx.wait();
-    });
-    
-    await Promise.all(txs);
-    console.log("Updated all keys!")
-  } else {
-    for (let i = 0; i < publicKeyParts.length; i++) {
-      const part = publicKeyParts[i];
-      const tx = await contract.editMailserverKey(domain, i, part);
-      await tx.wait();
-      console.log(`Updated key ${i}!`);
-    }
-  }
-}
-
-async function  testSelector(domain: string, selector: string) {
-  try {
-    const publicKeyParts = await formatDkimKey(domain, selector, false);
-    if (publicKeyParts) {
-      console.log(`Domain: ${domain}, Selector: ${selector} - Match found`);
-      return { match: true, selector: selector, domain: domain };
-    } else {
-      // console.log(`Domain: ${domain}, Selector: ${selector} - No match found`);
-    }
-  } catch (error) {
-    console.error(`Error processing domain: ${domain}, Selector: ${selector} - ${error}`);
-  }
-  return { match: false, selector: selector, domain: domain };
-}
-
-// Filename is a file where each line is a domain
-// This searches for default selectors like "google" or "default"
-async function getSelectors(filename: string, update_contract = false) {
-  const fs = require('fs');
-  // const selectors = ['google']
-  const selectors = ['google', 'default', 'mail', 'smtpapi', 'dkim', '200608', '20230601', '20221208', '20210112', 'v1', 'v2', 'v3', 'k1', 'k2', 'k3', 'hs1', 'hs2', 's1', 's2', 's3', 'sig1', 'sig2', 'sig3', 'selector', 'selector1', 'selector2', 'mindbox', 'bk', 'sm1', 'sm2', 'gmail', '10dkim1', '11dkim1', '12dkim1', 'memdkim', 'm1', 'mx', 'sel1', 'bk', 'scph1220', 'ml', 'pps1', 'scph0819', 'skiff1', 's1024', 'selector1'];
-
-  const data = fs.readFileSync(filename, 'utf8');
-  const domains = data.split('\n');
-  
-  let results = [];
-  let domainIndex = 0;
-  for (let domain of domains) {
-    const promises = [];
-    for (let selector of selectors) {
-      promises.push(testSelector(domain, selector));
-    }
-    domainIndex++;
-    results.push(...await Promise.all(promises));
-  }
-  const matchedDomains = new Set();
-  const matchedSelectors: {[key: string]: string[]} = {};
-  fs.writeFileSync('full_results.txt', JSON.stringify(results, null, 2));
-  for (let result of results) {
-    if(result.match) {
-      matchedDomains.add(result.domain);
-      if (!matchedSelectors[result.domain]) {
-        matchedSelectors[result.domain] = [];
-      }
-      matchedSelectors[result.domain].push(result.selector);
-    }
-  }
-  console.log("Domains with at least one matched selector: ");
-  console.log(Array.from(matchedDomains));
-
-  // Update mailserver contract with found keys
-  if (update_contract) {
-    for (let domain of Object.keys(matchedSelectors)) {
-      console.log(`Domain: ${domain}, Selectors: ${matchedSelectors[domain]}`);
-      for (let selector of matchedSelectors[domain]) {
-        await updateMailserverKeys(domain, selector, mailserver_address, default_abi, true);
-      }
-    }
-  }
-
-  fs.writeFileSync('domain_results.txt', JSON.stringify(Array.from(matchedDomains), null, 2));
-  fs.writeFileSync('selector_results.txt', JSON.stringify(matchedSelectors, null, 2));
-}
-
-let domain = process.argv[2] || 'gmail.com';
-let selector = process.argv[3] || '20230601';
-domain = 'protonmail.com';
-selector = 'protonmail3';
-// domain = 'pm.me';
-// selector = 'protonmail3';
-// updateMailserverKeys(domain, selector, mailserver_address, default_abi);
-getSelectors('src/dkim/domains.txt', true);
--- a/packages/scripts/.env.sample
+++ b/packages/scripts/.env.sample
@@ -0,0 +1,4 @@
+CHAIN_ID=
+INFURA_KEY=
+PRIVATE_KEY=
+DKIM_REGISTRY=
--- a/packages/helpers/src/dkim/domains.txt
+++ b/packages/helpers/src/dkim/domains.txt
--- a/packages/scripts/package.json
+++ b/packages/scripts/package.json
@@ -0,0 +1,21 @@
+{
+  "name": "@zk-email/scripts",
+  "version": "1.0.0",
+  "description": "",
+  "main": "index.js",
+  "scripts": {
+    "update-dkim-registry": "ts-node update-dkim-registry"
+  },
+  "keywords": [],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "circomlibjs": "^0.1.7",
+    "dotenv": "^16.3.1",
+    "ts-node": "^10.9.1",
+    "typescript": "^5.2.2"
+  },
+  "devDependencies": {
+    "@types/circomlibjs": "^0.1.5"
+  }
+}
--- a/packages/scripts/update-dkim-registry.ts
+++ b/packages/scripts/update-dkim-registry.ts
@@ -0,0 +1,267 @@
+import { ethers, JsonRpcProvider } from "ethers";
+import { buildPoseidon } from "circomlibjs";
+import dns from "dns";
+import forge from "node-forge";
+import { bigIntToChunkedBytes } from "@zk-email/helpers/src/binaryFormat";
+const fs = require("fs");
+import { abi } from "@zk-email/contracts/out/DKIMRegistry.sol/DKIMRegistry.json";
+require("dotenv").config();
+
+async function updateContract(domain: string, pubkeyHashes: string[]) {
+  if (!pubkeyHashes.length) {
+    return;
+  }
+
+  if (!process.env.PRIVATE_KEY) throw new Error("Env private key found");
+  if (!process.env.RPC_URL) throw new Error("Env RPC URL found");
+  if (!process.env.DKIM_REGISTRY) throw new Error("Env DKIM_REGISTRY found");
+
+  const provider = new JsonRpcProvider(process.env.RPC_URL);
+  const wallet = new ethers.Wallet(process.env.PRIVATE_KEY, provider);
+  const contract = new ethers.Contract(process.env.DKIM_REGISTRY, abi, wallet);
+
+  const hashes = pubkeyHashes.map((hash) => BigInt(hash));
+  const tx = await contract.setDKIMPublicKeyHashes(domain, hashes);
+  await tx.wait();
+
+  console.log(`Updated hashes for domain ${domain}. Tx: ${tx.hash}`);
+}
+
+async function getPublicKeyForDomainAndSelector(
+  domain: string,
+  selector: string,
+  print: boolean = true
+) {
+  // Construct the DKIM record name
+  let dkimRecordName = `${selector}._domainkey.${domain}`;
+  if (print) console.log(dkimRecordName);
+  // Lookup the DKIM record in DNS
+  let records;
+  try {
+    records = await dns.promises.resolveTxt(dkimRecordName);
+  } catch (err) {
+    if (print) console.error(err);
+    return;
+  }
+
+  if (!records.length) {
+    return;
+  }
+
+  // The DKIM record is a TXT record containing a string
+  // We need to parse this string to get the public key
+  let dkimRecord = records[0].join("");
+  let match = dkimRecord.match(/p=([^;]+)/);
+  if (!match) {
+    console.error("No public key found in DKIM record");
+    return;
+  }
+
+  // The public key is base64 encoded, we need to decode it
+  let pubkey = match[1];
+  let binaryKey = Buffer.from(pubkey, "base64").toString("base64");
+
+  // Get match
+  let matches = binaryKey.match(/.{1,64}/g);
+  if (!matches) {
+    console.error("No matches found");
+    return;
+  }
+  let formattedKey = matches.join("\n");
+  if (print) console.log("Key: ", formattedKey);
+
+  // Convert to PEM format
+  let pemKey = `-----BEGIN PUBLIC KEY-----\n${formattedKey}\n-----END PUBLIC KEY-----`;
+
+  // Parse the RSA public key
+  let publicKey = forge.pki.publicKeyFromPem(pemKey);
+
+  // Get the modulus n only
+  let n = publicKey.n;
+  if (print) console.log("Modulus n:", n.toString(16));
+
+  return BigInt(publicKey.n.toString());
+}
+
+async function checkSelector(domain: string, selector: string) {
+  try {
+    const publicKey = await getPublicKeyForDomainAndSelector(
+      domain,
+      selector,
+      false
+    );
+    if (publicKey) {
+      console.log(`Domain: ${domain}, Selector: ${selector} - Match found`);
+      return {
+        match: true,
+        selector: selector,
+        domain: domain,
+        publicKey,
+      };
+    } else {
+      // console.log(`Domain: ${domain}, Selector: ${selector} - No match found`);
+    }
+  } catch (error) {
+    console.error(
+      `Error processing domain: ${domain}, Selector: ${selector} - ${error}`
+    );
+  }
+
+  return {
+    match: false,
+    selector: selector,
+    domain: domain,
+    publicKey: null,
+  };
+}
+
+// Filename is a file where each line is a domain
+// This searches for default selectors like "google" or "default"
+async function getDKIMPublicKeysForDomains(filename: string) {
+  const domains = fs.readFileSync(filename, "utf8").split("\n");
+  const selectors = [
+    "google",
+    "default",
+    "mail",
+    "smtpapi",
+    "dkim",
+    "200608",
+    "20230601",
+    "20221208",
+    "20210112",
+    "v1",
+    "v2",
+    "v3",
+    "k1",
+    "k2",
+    "k3",
+    "hs1",
+    "hs2",
+    "s1",
+    "s2",
+    "s3",
+    "sig1",
+    "sig2",
+    "sig3",
+    "selector",
+    "selector1",
+    "selector2",
+    "mindbox",
+    "bk",
+    "sm1",
+    "sm2",
+    "gmail",
+    "10dkim1",
+    "11dkim1",
+    "12dkim1",
+    "memdkim",
+    "m1",
+    "mx",
+    "sel1",
+    "bk",
+    "scph1220",
+    "ml",
+    "pps1",
+    "scph0819",
+    "skiff1",
+    "s1024",
+    "selector1",
+  ];
+
+  let results = [];
+
+  for (let domain of domains) {
+    const promises = [];
+    for (let selector of selectors) {
+      promises.push(checkSelector(domain, selector));
+    }
+    results.push(...(await Promise.all(promises)));
+  }
+
+  const matchedSelectors: { [key: string]: string[] } = {};
+
+  for (let result of results) {
+    if (result.match && result.publicKey) {
+      if (!matchedSelectors[result.domain]) {
+        matchedSelectors[result.domain] = [];
+      }
+
+      const publicKey = result.publicKey.toString();
+
+      if (!matchedSelectors[result.domain].includes(publicKey)) {
+        matchedSelectors[result.domain].push(publicKey);
+      }
+    }
+  }
+
+  return matchedSelectors;
+}
+
+async function updateDKIMRegistry(
+  { writeToFile } = {
+    writeToFile: false,
+  }
+) {
+  const domainsFile = "./domains.txt";
+  const domainPubKeyMap = await getDKIMPublicKeysForDomains(domainsFile);
+
+  if (writeToFile) {
+    fs.writeFileSync(
+      "out/domain-dkim-keys.json",
+      JSON.stringify(domainPubKeyMap, null, 2)
+    );
+  }
+
+  // const domainPubKeyMap = JSON.parse(
+  //   fs.readFileSync("out/domain-dkim-keys.json").toString()
+  // );
+
+  // Saving pubkeys into chunks of 121 * 17
+  // This is what is used in EmailVerifier.cicrom
+  // Can be used at https://zkrepl.dev/?gist=43ce7dce2466c63812f6efec5b13aa73 to get pubkey hash
+  const chunkedDKIMPubKeyMap: { [key: string]: string[][] } = {};
+  for (let domain of Object.keys(domainPubKeyMap)) {
+    for (let publicKey of domainPubKeyMap[domain]) {
+      const pubkeyChunked = bigIntToChunkedBytes(BigInt(publicKey), 121, 17);
+
+      if (!chunkedDKIMPubKeyMap[domain]) {
+        chunkedDKIMPubKeyMap[domain] = [];
+      }
+      chunkedDKIMPubKeyMap[domain].push(pubkeyChunked.map((s) => s.toString()));
+    }
+  }
+  if (writeToFile) {
+    fs.writeFileSync(
+      "out/domain-dkim-keys-chunked.json",
+      JSON.stringify(chunkedDKIMPubKeyMap, null, 2)
+    );
+  }
+
+  // Generate pub key hash using 242 * 9 chunks (Poseidon lib don't take more than 16 inputs)
+  const domainHashedPubKeyMap: { [key: string]: string[] } = {};
+  const poseidon = await buildPoseidon();
+  for (let domain of Object.keys(domainPubKeyMap)) {
+    for (let publicKey of domainPubKeyMap[domain]) {
+      const pubkeyChunked = bigIntToChunkedBytes(BigInt(publicKey), 242, 9);
+      const hash = poseidon(pubkeyChunked);
+
+      if (!domainHashedPubKeyMap[domain]) {
+        domainHashedPubKeyMap[domain] = [];
+      }
+      domainHashedPubKeyMap[domain].push(poseidon.F.toObject(hash).toString());
+    }
+  }
+  if (writeToFile) {
+    fs.writeFileSync(
+      "out/domain-dkim-key-hashes.json",
+      JSON.stringify(domainHashedPubKeyMap, null, 2)
+    );
+  }
+
+  // Update Mailserver contract with found keys
+  for (let domain of Object.keys(domainHashedPubKeyMap)) {
+    await updateContract(domain, domainHashedPubKeyMap[domain]);
+  }
+}
+
+updateDKIMRegistry({ writeToFile: true });
--- a/yarn.lock
+++ b/yarn.lock
@@ -3706,6 +3706,13 @@ __metadata:
  languageName: node
  linkType: hard

+"@types/circomlibjs@npm:^0.1.5":
+  version: 0.1.5
+  resolution: "@types/circomlibjs@npm:0.1.5"
+  checksum: 5f9ad10f2b3047f31c915a814984a1a531fd02595d085b9046f66b4526123069b199e8e69c60abdf6ec4b50f509aba902567d5dc545210acd3a322df31d8dcfb
+  languageName: node
+  linkType: hard
+
 "@types/connect@npm:^3.4.33":
  version: 3.4.35
  resolution: "@types/connect@npm:3.4.35"
@@ -4716,6 +4723,7 @@ __metadata:
  resolution: "@zk-email/contracts@workspace:packages/contracts"
  dependencies:
    "@openzeppelin/contracts": ^4.9.3
+    dotenv: ^16.3.1
  languageName: unknown
  linkType: soft

@@ -4757,6 +4765,18 @@ __metadata:
  languageName: unknown
  linkType: soft

+"@zk-email/scripts@workspace:packages/scripts":
+  version: 0.0.0-use.local
+  resolution: "@zk-email/scripts@workspace:packages/scripts"
+  dependencies:
+    "@types/circomlibjs": ^0.1.5
+    circomlibjs: ^0.1.7
+    dotenv: ^16.3.1
+    ts-node: ^10.9.1
+    typescript: ^5.2.2
+  languageName: unknown
+  linkType: soft
+
 "@zk-email/twitter-verifier-contracts@workspace:packages/twitter-verifier-contracts":
  version: 0.0.0-use.local
  resolution: "@zk-email/twitter-verifier-contracts@workspace:packages/twitter-verifier-contracts"