Generate a MultiversX wallet in a specific shard

generate-wallet.ts

import { bech32 } from 'bech32';
import { generateMnemonic, mnemonicToSeedSync, validateMnemonic } from 'bip39';
import { derivePath } from 'ed25519-hd-key';
import * as tweetnacl from 'tweetnacl';

const MNEMONIC_LENGTH = 256;
const HD_PREFIX = "m/44'/508'/0'/0'";
const ERD = "erd";
const HD_PATH = `${HD_PREFIX}/0'`;

interface GenerateWalletResponse {
  mnemonic: string;
  address: string;
}

const generateMnemonicInShard = (expectedShard: number): Promise<GenerateWalletResponse> =>
  new Promise((resolve, reject) => {
    while (true) {
      const mnemonic = generateMnemonic(MNEMONIC_LENGTH);
      const { publicKey, secretKey } = keysFromMnemonic(mnemonic);

      const shardID = computeShardID(Buffer.from(secretKey));

      if (shardID !== expectedShard) {
        continue;
      }

      resolve({
        mnemonic,
        address: addressFromHexPublicKey(publicKey),
      });

      break;
    }
  });

const addressFromHexPublicKey = (publicKey: Uint8Array) => {
  let words = bech32.toWords(publicKey);
  return bech32.encode(ERD, words);
};

const keysFromMnemonic = (mnemonic: string) => {
  if (!validateMnemonic(mnemonic)) {
    throw new Error("Invalid mnemonic format");
  }
  const seed = mnemonicToSeedSync(mnemonic);
  const { key } = derivePath(HD_PATH, seed.toString("hex"));
  const privateKey = Uint8Array.from(key);
  const kp = tweetnacl.sign.keyPair.fromSeed(privateKey);

  return {
    publicKey: kp.publicKey,
    secretKey: kp.secretKey,
  };
};

const computeShardID = (pubKey: Buffer) => {
  const startingIndex = pubKey.length - 1;

  const usedBuffer = pubKey.slice(startingIndex);

  let addr = 0;
  for (let i = 0; i < usedBuffer.length; i++) {
    addr = (addr << 8) + usedBuffer[i];
  }

  let n = Math.ceil(Math.log2(3));
  let maskHigh = (1 << n) - 1;
  let maskLow = (1 << (n - 1)) - 1;

  let shard = addr & maskHigh;
  if (shard > 2) {
    shard = addr & maskLow;
  }

  return shard;
};

generateMnemonicInShard(2).then(console.log);