AutistMask/src/popup/views/addWallet.js

const { $, showView, showFlash } = require("./helpers");
const {
    generateMnemonic,
    hdWalletFromMnemonic,
    isValidMnemonic,
    addressFromPrivateKey,
    hdWalletFromXprv,
    isValidXprv,
} = require("../../shared/wallet");
const { encryptWithPassword } = require("../../shared/vault");
const { state, saveState } = require("../../shared/state");
const { scanForAddresses } = require("../../shared/balances");

let currentMode = "mnemonic";

const MODES = ["mnemonic", "privkey", "xprv"];

const PASSWORD_HINTS = {
    mnemonic:
        "This password encrypts your recovery phrase on this device. You will need it to send funds.",
    privkey:
        "This password encrypts your private key on this device. You will need it to send funds.",
    xprv: "This password encrypts your key on this device. You will need it to send funds.",
};

function switchMode(mode) {
    currentMode = mode;
    for (const m of MODES) {
        $("add-wallet-section-" + m).classList.toggle("hidden", m !== mode);
        const tab = $("tab-" + m);
        tab.classList.toggle("border-fg", m === mode);
        tab.classList.toggle("font-bold", m === mode);
        tab.classList.toggle("border-transparent", m !== mode);
        tab.classList.toggle("text-muted", m !== mode);
    }
    $("add-wallet-password-hint").textContent = PASSWORD_HINTS[mode];
}

function show() {
    $("wallet-mnemonic").value = "";
    $("import-private-key").value = "";
    $("import-xprv-key").value = "";
    $("add-wallet-password").value = "";
    $("add-wallet-password-confirm").value = "";
    $("add-wallet-phrase-warning").classList.add("hidden");
    switchMode("mnemonic");
    showView("add-wallet");
}

function validatePassword() {
    const pw = $("add-wallet-password").value;
    const pw2 = $("add-wallet-password-confirm").value;
    if (!pw) {
        showFlash("Please choose a password.");
        return null;
    }
    if (pw.length < 12) {
        showFlash("Password must be at least 12 characters.");
        return null;
    }
    if (pw !== pw2) {
        showFlash("Passwords do not match.");
        return null;
    }
    return pw;
}

async function importMnemonic(ctx) {
    const mnemonic = $("wallet-mnemonic").value.trim();
    if (!mnemonic) {
        showFlash("Enter a recovery phrase or press the die to generate one.");
        return;
    }
    const words = mnemonic.split(/\s+/);
    if (words.length !== 12 && words.length !== 24) {
        showFlash(
            "Recovery phrase must be 12 or 24 words. You entered " +
                words.length +
                ".",
        );
        return;
    }
    if (!isValidMnemonic(mnemonic)) {
        showFlash("Invalid recovery phrase. Check for typos.");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const { xpub, firstAddress } = hdWalletFromMnemonic(mnemonic);
    const duplicate = state.wallets.find(
        (w) =>
            w.type === "hd" &&
            w.addresses[0] &&
            w.addresses[0].address.toLowerCase() === firstAddress.toLowerCase(),
    );
    if (duplicate) {
        showFlash(
            "This recovery phrase is already added (" + duplicate.name + ").",
        );
        return;
    }
    const encrypted = await encryptWithPassword(mnemonic, pw);
    const walletNum = state.wallets.length + 1;
    const wallet = {
        type: "hd",
        name: "Wallet " + walletNum,
        xpub: xpub,
        encryptedSecret: encrypted,
        nextIndex: 1,
        addresses: [
            { address: firstAddress, balance: "0.0000", tokenBalances: [] },
        ],
    };
    state.wallets.push(wallet);
    state.hasWallet = true;
    await saveState();
    ctx.renderWalletList();
    showView("main");

    // Scan for used HD addresses beyond index 0.
    showFlash("Scanning for addresses...", 30000);
    const scan = await scanForAddresses(xpub, state.rpcUrl);
    if (scan.addresses.length > 1) {
        wallet.addresses = scan.addresses.map((a) => ({
            address: a.address,
            balance: "0.0000",
            tokenBalances: [],
        }));
        wallet.nextIndex = scan.nextIndex;
        await saveState();
        ctx.renderWalletList();
        showFlash("Found " + scan.addresses.length + " addresses.");
    } else {
        showFlash("Ready.", 1000);
    }

    ctx.doRefreshAndRender();
}

async function importPrivateKey(ctx) {
    const key = $("import-private-key").value.trim();
    if (!key) {
        showFlash("Please enter your private key.");
        return;
    }
    let addr;
    try {
        addr = addressFromPrivateKey(key);
    } catch (e) {
        showFlash("Invalid private key.");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const duplicate = state.wallets.find(
        (w) =>
            w.type === "key" &&
            w.addresses[0] &&
            w.addresses[0].address.toLowerCase() === addr.toLowerCase(),
    );
    if (duplicate) {
        showFlash(
            "This private key is already added (" + duplicate.name + ").",
        );
        return;
    }
    const encrypted = await encryptWithPassword(key, pw);
    const walletNum = state.wallets.length + 1;
    state.wallets.push({
        type: "key",
        name: "Wallet " + walletNum,
        encryptedSecret: encrypted,
        addresses: [{ address: addr, balance: "0.0000", tokenBalances: [] }],
    });
    state.hasWallet = true;
    await saveState();
    ctx.renderWalletList();
    showView("main");

    ctx.doRefreshAndRender();
}

async function importXprvKey(ctx) {
    const xprv = $("import-xprv-key").value.trim();
    if (!xprv) {
        showFlash("Please enter your extended private key.");
        return;
    }
    if (!isValidXprv(xprv)) {
        showFlash("Invalid extended private key.");
        return;
    }
    let result;
    try {
        result = hdWalletFromXprv(xprv);
    } catch (e) {
        showFlash("Invalid extended private key.");
        return;
    }
    const { xpub, firstAddress } = result;
    const duplicate = state.wallets.find(
        (w) =>
            (w.type === "hd" || w.type === "xprv") &&
            w.addresses[0] &&
            w.addresses[0].address.toLowerCase() === firstAddress.toLowerCase(),
    );
    if (duplicate) {
        showFlash("This key is already added (" + duplicate.name + ").");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const encrypted = await encryptWithPassword(xprv, pw);
    const walletNum = state.wallets.length + 1;
    const wallet = {
        type: "xprv",
        name: "Wallet " + walletNum,
        xpub: xpub,
        encryptedSecret: encrypted,
        nextIndex: 1,
        addresses: [
            { address: firstAddress, balance: "0.0000", tokenBalances: [] },
        ],
    };
    state.wallets.push(wallet);
    state.hasWallet = true;
    await saveState();
    ctx.renderWalletList();
    showView("main");

    // Scan for used HD addresses beyond index 0.
    showFlash("Scanning for addresses...", 30000);
    const scan = await scanForAddresses(xpub, state.rpcUrl);
    if (scan.addresses.length > 1) {
        wallet.addresses = scan.addresses.map((a) => ({
            address: a.address,
            balance: "0.0000",
            tokenBalances: [],
        }));
        wallet.nextIndex = scan.nextIndex;
        await saveState();
        ctx.renderWalletList();
        showFlash("Found " + scan.addresses.length + " addresses.");
    } else {
        showFlash("Ready.", 1000);
    }

    ctx.doRefreshAndRender();
}

function init(ctx) {
    // Tab click handlers
    $("tab-mnemonic").addEventListener("click", () => switchMode("mnemonic"));
    $("tab-privkey").addEventListener("click", () => switchMode("privkey"));
    $("tab-xprv").addEventListener("click", () => switchMode("xprv"));

    // Generate mnemonic
    $("btn-generate-phrase").addEventListener("click", () => {
        $("wallet-mnemonic").value = generateMnemonic();
        $("add-wallet-phrase-warning").classList.remove("hidden");
    });

    // Import / confirm
    $("btn-add-wallet-confirm").addEventListener("click", async () => {
        if (currentMode === "mnemonic") {
            await importMnemonic(ctx);
        } else if (currentMode === "privkey") {
            await importPrivateKey(ctx);
        } else if (currentMode === "xprv") {
            await importXprvKey(ctx);
        }
    });

    // Back button
    $("btn-add-wallet-back").addEventListener("click", () => {
        if (!state.hasWallet) {
            showView("welcome");
        } else {
            ctx.renderWalletList();
            showView("main");
        }
    });
}

module.exports = { init, show };