diff --git a/src/cli.rs b/src/cli.rs index 824322b..2004d17 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -2,9 +2,8 @@ use clap::{Parser, Subcommand}; use fake::{faker::lorem::en::Sentence, Fake}; use std::io; mod commands; -use self::commands::{movejs, proofs}; use crate::utils::{command_runner::RealCommandRunner, filesystem_operations::RealFileSystemOps}; -use commands::{circuit, compile, setup, verifier}; +use commands::{all, circuit, compile, movejs, proofs, setup, token, verifier}; /// Represents the command line interface for the Zero Knowledge Whitelist Tool. /// Deriving `Parser` from clap allows for automatic parsing of command line arguments. @@ -35,8 +34,12 @@ pub enum SubCommand { Verifier, /// Moves the contents of `circuit_js` on parent directory Movejs, + /// Generates a sample token solidity contract, to be used together with verifier. + Token, /// Generates proofs using an input file, with a default value of "addresses.txt". Proofs(ProofsCommand), + /// Run all the commands one after the other, {circuit, compile, setup, verifier, movejs, proofs} using an input file, with a default value of "addresses.txt" + All(AllCommand), } #[derive(Parser, PartialEq, Debug)] @@ -45,6 +48,12 @@ pub struct ProofsCommand { pub input_file: String, } +#[derive(Parser, PartialEq, Debug)] +pub struct AllCommand { + #[clap(long, default_value = "addresses.txt")] + pub input_file: String, +} + /// The entry point of the application. /// Parses command line arguments and executes the corresponding subcommand. pub fn run_cli() -> std::io::Result<()> { @@ -64,6 +73,16 @@ pub fn run_cli() -> std::io::Result<()> { proofs::handle_proofs_subcommand(&runner, &proofs_command.input_file, &file_system_ops) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))? } + SubCommand::Token => token::handle_token_subcommand()?, + SubCommand::All(all_command) => { + all::handle_all_command( + runner, + random_name, + random_text, + file_system_ops, + all_command, + )?; + } }; Ok(()) @@ -124,4 +143,26 @@ mod tests { }) ); } + + #[test] + fn test_parse_all_subcommand_with_default_value() { + let args = Cli::parse_from(&["zk_whitelist", "all"]); + assert_eq!( + args.subcmd, + SubCommand::All(AllCommand { + input_file: "addresses.txt".to_string() + }) + ); + } + + #[test] + fn test_parse_all_subcommand_with_custom_value() { + let args = Cli::parse_from(&["zk_whitelist", "all", "--input-file", "custom.txt"]); + assert_eq!( + args.subcmd, + SubCommand::All(AllCommand { + input_file: "custom.txt".to_string() + }) + ); + } } diff --git a/src/cli/commands/all.rs b/src/cli/commands/all.rs new file mode 100644 index 0000000..7c56fec --- /dev/null +++ b/src/cli/commands/all.rs @@ -0,0 +1,25 @@ +use crate::{ + cli::AllCommand, + utils::{command_runner::RealCommandRunner, filesystem_operations::RealFileSystemOps}, +}; +use std::io; + +use super::{circuit, compile, movejs, proofs, setup, token, verifier}; + +pub fn handle_all_command( + runner: RealCommandRunner, + random_name: String, + random_text: String, + file_system_ops: RealFileSystemOps, + all_command: AllCommand, +) -> Result<(), io::Error> { + circuit::handle_circuit_subcommand()?; + compile::handle_compile_subcommand(&runner)?; + setup::handle_setup_subcommand(&runner, random_name.clone(), random_text.clone())?; + verifier::handle_verifier_subcommand(&runner)?; + token::handle_token_subcommand()?; + movejs::handle_movejs_subcommand(&file_system_ops)?; + proofs::handle_proofs_subcommand(&runner, &all_command.input_file, &file_system_ops) + .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?; + Ok(()) +} diff --git a/src/cli/commands/mod.rs b/src/cli/commands/mod.rs index 0a33b9a..081c108 100644 --- a/src/cli/commands/mod.rs +++ b/src/cli/commands/mod.rs @@ -1,6 +1,8 @@ +pub mod all; pub mod circuit; pub mod compile; pub mod movejs; pub mod proofs; pub mod setup; +pub mod token; pub mod verifier; diff --git a/src/cli/commands/token.rs b/src/cli/commands/token.rs new file mode 100644 index 0000000..6a164f6 --- /dev/null +++ b/src/cli/commands/token.rs @@ -0,0 +1,66 @@ +use std::env; +use std::fs::File; +use std::io::{self, Write}; +use std::path::Path; + +/// Copies a token file template to the current directory. +/// +/// This function is intended to be called when a verifier is available. +/// It provides a sample token solidity contract to be used with the verifier +/// +/// # Errors +/// Returns an `io::Result` wrapping any I/O error that occurs. +fn copy_token_file() -> io::Result<()> { + // Obtain the current working directory + let current_dir = env::current_dir()?; + // Construct a path for the new solidity contract + let circuit_path = Path::new(¤t_dir).join("zkToken.sol"); + // Create a new file at the constructed path + let mut file = File::create(circuit_path)?; + // Write the contents of the template file into the new file + file.write_all(include_bytes!("../../../templates/zkToken.sol"))?; + Ok(()) +} + +/// Handles the `circuit` CLI subcommand. +/// +/// This function acts as a handler for the `circuit` subcommand. +/// It calls the `copy_circuit_file` function to perform the actual work. +/// +/// # Returns +/// Returns an `io::Result` to indicate success or any I/O error that occurs. +pub fn handle_token_subcommand() -> std::io::Result<()> { + copy_token_file() +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::path::Path; + + /// Tests the functionality of the `handle_circuit_subcommand` function. + /// + /// This test ensures that the `handle_circuit_subcommand` function correctly + /// creates a new `circuit.circom` file in the current directory. + /// + /// # Returns + /// Returns an `io::Result` to indicate the success or failure of the test. + #[test] + fn test_handle_token_subcommand() -> std::io::Result<()> { + // Execute the function under test + handle_token_subcommand()?; + + // Obtain the current working directory + let current_dir = std::env::current_dir()?; + // Construct the path of the circuit file + let circuit_path = Path::new(¤t_dir).join("zkToken.sol"); + // Assert that the file has been created + assert!(circuit_path.exists()); + + // Clean up by removing the created file + fs::remove_file(circuit_path)?; + + Ok(()) + } +} diff --git a/templates/zkToken.sol b/templates/zkToken.sol new file mode 100644 index 0000000..4448215 --- /dev/null +++ b/templates/zkToken.sol @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.20; + +// import the verifier that the program created +import "./verifier.sol"; +import "@openzeppelin/contracts/token/ERC20/ERC20.sol"; +import "@openzeppelin/contracts/utils/Strings.sol"; + +/// @title ZKToken Contract +/// @notice This contract represents an ERC20 token with minting only for ZK proven accounts. +/// @notice Requires a verifier circuit contracts +contract ZKToken is ERC20 { + Groth16Verifier public verifier; + mapping(address => bool) public claimed; + + constructor() ERC20("YourToken", "YTK") { + verifier = new Groth16Verifier(); + } + + /* + * @notice Mints new tokens after verifying a provided proof. + * @param pA, pB, pC, pubSignals The ZK proofs from proofs file. + * @return A boolean value indicating whether the function executed successfully. Reverts otherwise. + */ + function mint(uint[2] calldata _pA, uint[2][2] calldata _pB, uint[2] calldata _pC, uint[2] calldata _pubSignals ) public returns (bool) { + // Convert msg.sender address to decimal + uint256 senderDecimalAddress = uint256(uint160(msg.sender)); + + // Ensure the proof is for sender + require(senderDecimalAddress == _pubSignals[1], "Not your proof or invalid input"); + + // Ensure the tokens haven't been claimed yet + require(!claimed[msg.sender], "Tokens already claimed"); + + // Verify the proof + require(verifier.verifyProof(_pA, _pB, _pC, _pubSignals), "Invalid proof"); + + // Mark as claimed and mint the tokens + claimed[msg.sender] = true; + _mint(msg.sender, 10 * 10 ** decimals()); + return true; + } +} \ No newline at end of file