add erc721

README_ZH.md

@@ -1,10 +1,10 @@
 # Assets Snapshot: Create ERC20, ERC721, and ERC1155 Token Snapshot
 
-一个命令行工具和SDK，用于创建任何ERC20、ERC721和ERC1155代币的快照，将数据同步到SQLite，并允许导出为JSON或CSV格式。
+一个命令行工具和 SDK，用于创建任何 ERC20、ERC721 和 ERC1155 代币的快照，将数据同步到 SQLite，并允许导出为 JSON 或 CSV 格式。
 
-- 无需本地Ethereum节点。
+- 无需本地 Ethereum 节点。
 - 在失败时会自动恢复。
-- 经测试与Alchemy兼容。
+- 经测试与 Alchemy 兼容。
 
 ## Getting Started
 
@@ -44,7 +44,7 @@ assets-snapshot
 
 ### contractAddress
 
-您的ERC20、ERC721和ERC1155合约地址。
+您的 ERC20、ERC721 和 ERC1155 合约地址。
 
 ### fromBlock
 

src/app.ts

@@ -1,9 +1,9 @@
 import { formatUnits } from "viem";
 import { getConfig } from "./config";
 import { getBlockNumber, getERC20Decimals, getERC20Transfer, getERC721Transfer } from "./contract";
-import { getAssetTransfersERC20, getMaxBlockNum, IAssetTransactionsERC20, insertTransactionsERC20 } from "./db";
+import { getAssetTransfersERC20, getAssetTransfersERC721, getMaxBlockNum, IAssetTransactionsERC20, IAssetTransactionsERC721, insertTransactionsERC20, insertTransactionsERC721 } from "./db";
 import { sleep } from "./helper";
-import { exportBalancesERC20, IExportBalancesERC20 } from "./export";
+import { exportBalancesERC20, exportBalancesERC721, IExportBalancesERC20, IExportBalancesERC721 } from "./export";
 
 export const start = async () => {
   const config = getConfig();
@@ -147,8 +147,133 @@ const handleERC20 = async () => {
 };
 
 const handleERC721 = async () => {
-  const logs = await getERC721Transfer(17572615n, 17572617n);
-  console.log(logs);
+  const config = getConfig();
+
+  let fromBlock = BigInt(config.fromBlock);
+  let toBlock = 0n;
+
+  const maxBlockNum = await getMaxBlockNum();
+
+  if (maxBlockNum) {
+    console.log("Resuming from the last downloaded block #", maxBlockNum);
+    fromBlock = BigInt(maxBlockNum) + 1n;
+  }
+
+  if (config.toBlock === "latest") {
+    toBlock = await getBlockNumber();
+  } else {
+    toBlock = BigInt(config.toBlock);
+  }
+
+  console.log("From %d to %d", fromBlock, toBlock);
+
+  const blocksPerBatch = BigInt(config.blocksPerBatch);
+  const delay = config.delay;
+
+  let start = fromBlock;
+  let end = fromBlock + blocksPerBatch;
+  let i = 0;
+
+  while (end < toBlock) {
+    i++;
+
+    if (delay) {
+      await sleep(delay);
+    }
+
+    console.log("Batch", i, " From", start, "to", end);
+
+    const logs = await getERC721Transfer(start, end);
+
+    const transactions = logs.map((log) => ({
+      blockNum: Number(log.blockNum),
+      hash: log.hash,
+      sender: log.from,
+      recipient: log.to,
+      logIndex: log.logIndex,
+      tokenId: log.tokenId.toString()
+    }));
+
+    console.log("Transactions count ", transactions.length);
+
+    await insertTransactionsERC721(transactions);
+
+    start = end + 1n;
+    end = start + blocksPerBatch;
+
+    if (end > toBlock) {
+      end = toBlock;
+    }
+  }
+
+  console.log("Calculating balances of %s ", config.name);
+
+  const balances = new Map<string, { deposits: string[]; withdrawals: string[] }>();
+  const closingBalances: IExportBalancesERC721[] = [];
+
+  const setDeposits = (event: IAssetTransactionsERC721) => {
+    const wallet = event.recipient;
+
+    let deposits = (balances.get(wallet) || {}).deposits || [];
+    let withdrawals = (balances.get(wallet) || {}).withdrawals || [];
+
+    if (event.tokenId) {
+      deposits = [...deposits, event.tokenId];
+      balances.set(wallet, { deposits, withdrawals });
+    } else {
+      throw new TypeError("invalid tokenId value");
+    }
+  };
+
+  const setWithdrawals = (event: IAssetTransactionsERC721) => {
+    const wallet = event.sender;
+
+    let deposits = (balances.get(wallet) || {}).deposits || [];
+    let withdrawals = (balances.get(wallet) || {}).withdrawals || [];
+
+    if (event.tokenId) {
+      withdrawals = [...withdrawals, event.tokenId];
+      balances.set(wallet, { deposits, withdrawals });
+    } else {
+      throw new TypeError("invalid tokenId value");
+    }
+  };
+
+  let isComplete = false;
+  let lastId = 0;
+
+  while (!isComplete) {
+    const transactions = await getAssetTransfersERC721(lastId);
+    if (transactions && transactions.length > 0) {
+      lastId = transactions[transactions.length - 1].id;
+
+      for (const event of transactions) {
+        setDeposits(event);
+        setWithdrawals(event);
+      }
+    } else {
+      isComplete = true;
+    }
+  }
+
+  for (const [key, value] of balances.entries()) {
+    if (key === "0x0000000000000000000000000000000000000000") continue;
+
+    const tokenIds = value.deposits.filter((x) => !value.withdrawals.includes(x));
+
+    closingBalances.push({
+      wallet: key,
+      tokenIds
+    });
+  }
+
+  const filterBalances = closingBalances.filter((b) => b.tokenIds.length > 0);
+
+  console.log("Exporting balances");
+  await exportBalancesERC721(filterBalances);
+  console.log("Exporting balances complete");
 };
 
-const handleERC1155 = async () => {};
+const handleERC1155 = async () => {
+  console.warn("Not implemented yet");
+};

src/db.ts

@@ -147,3 +147,79 @@ export const getAssetTransfersERC20 = async (lastId: number = 0) => {
     });
   });
 };
+
+export interface IInsertTransactionsERC721 {
+  blockNum: number;
+  hash: string;
+  sender: string;
+  recipient: string;
+  tokenId: string;
+  logIndex: number;
+}
+
+export const insertTransactionsERC721 = async (transactions: IInsertTransactionsERC721[]) => {
+  const config = getConfig();
+
+  if (config.category !== "ERC721") {
+    throw new Error("Invalid category");
+  }
+
+  const tableName = `transaction_${config.name}`;
+  const sql = `
+    INSERT INTO ${tableName} (blockNum, hash, sender, recipient, tokenId, logIndex)
+    VALUES (?, ?, ?, ?, ?, ?)
+  `;
+
+  return new Promise<void>((resolve, reject) => {
+    db.serialize(() => {
+      db.run("BEGIN TRANSACTION");
+      const stmt = db.prepare(sql);
+
+      try {
+        for (const tx of transactions) {
+          stmt.run(tx.blockNum, tx.hash, tx.sender, tx.recipient, tx.tokenId, tx.logIndex);
+        }
+
+        stmt.finalize();
+        db.run("COMMIT", (err) => {
+          if (err) reject(err);
+          resolve();
+        });
+      } catch (error) {
+        db.run("ROLLBACK", () => reject(error));
+      }
+    });
+  });
+};
+
+export interface IAssetTransactionsERC721 {
+  id: number;
+  sender: string;
+  recipient: string;
+  tokenId: string;
+}
+
+export const getAssetTransfersERC721 = async (lastId: number = 0) => {
+  const config = getConfig();
+
+  if (config.category !== "ERC721") {
+    throw new Error("Invalid category");
+  }
+
+  const tableName = `transaction_${config.name}`;
+
+  const sql = `
+  SELECT id, sender, recipient, tokenId
+  FROM ${tableName}
+  WHERE id > ?
+  ORDER BY id ASC
+  LIMIT ?
+`;
+
+  return new Promise<IAssetTransactionsERC721[] | null>((resolve, reject) => {
+    db.all(sql, [lastId, config.blocksPerBatch], (err, rows) => {
+      if (err) reject(err);
+      resolve(rows as IAssetTransactionsERC721[] | null);
+    });
+  });
+};

src/export.ts

@@ -15,7 +15,47 @@ export const exportBalancesERC20 = (balances: IExportBalancesERC20[]) => {
       reject("Invalid category");
     }
 
-    const rootDir = path.resolve(__dirname, "../"); // 根据你的项目结构调整
+    const rootDir = path.resolve(__dirname, "../"); // Adapt to your project structure
+
+    const balancesDir = path.join(rootDir, "balances");
+
+    if (!fs.existsSync(balancesDir)) {
+      fs.mkdirSync(balancesDir);
+    }
+
+    if (config.format === "json") {
+      const filePath = path.join(balancesDir, `${config.name}.json`);
+      fs.writeFileSync(filePath, JSON.stringify(balances, null, 2), "utf-8");
+      resolve(true);
+    } else if (config.format === "csv") {
+      const ws = fs.createWriteStream(path.join(balancesDir, `${config.name}.csv`));
+
+      fastcsv
+        .write(balances, { headers: true })
+        .pipe(ws)
+        .on("finish", () => {
+          console.log("CSV file was written successfully");
+          resolve(true);
+        });
+    } else {
+      reject("Invalid format");
+    }
+  });
+};
+
+export interface IExportBalancesERC721 {
+  wallet: string;
+  tokenIds: string[];
+}
+
+export const exportBalancesERC721 = (balances: IExportBalancesERC721[]) => {
+  return new Promise((resolve, reject) => {
+    const config = getConfig();
+    if (config.category !== "ERC721") {
+      reject("Invalid category");
+    }
+
+    const rootDir = path.resolve(__dirname, "../"); // Adapt to your project structure
 
     const balancesDir = path.join(rootDir, "balances");