From b60cf7d5e77ce1431bc8f8dd9cd2493acc0adeb7 Mon Sep 17 00:00:00 2001
From: lonerapier <lonerapier@proton.me>
Date: Wed, 11 Sep 2024 21:07:49 +0530
Subject: [PATCH] pending work

---
 circuits/test/http/codegen.test.ts            | 16 ++--
 docs/http.md                                  |  0
 docs/pabuild.md                               | 44 ++++++----
 .../{extractor => lockfile}/two_keys.json     |  0
 .../value_array_nested.json                   |  0
 .../value_array_number.json                   |  0
 .../value_array_object.json                   |  0
 .../value_array_string.json                   |  0
 .../{extractor => lockfile}/value_number.json |  0
 .../{extractor => lockfile}/value_object.json |  0
 .../{extractor => lockfile}/value_string.json |  0
 src/json.rs                                   | 25 ++++--
 src/main.rs                                   | 30 ++++++-
 src/witness.rs                                | 86 ++++++++++++++++---
 14 files changed, 153 insertions(+), 48 deletions(-)
 create mode 100644 docs/http.md
 rename examples/json/{extractor => lockfile}/two_keys.json (100%)
 rename examples/json/{extractor => lockfile}/value_array_nested.json (100%)
 rename examples/json/{extractor => lockfile}/value_array_number.json (100%)
 rename examples/json/{extractor => lockfile}/value_array_object.json (100%)
 rename examples/json/{extractor => lockfile}/value_array_string.json (100%)
 rename examples/json/{extractor => lockfile}/value_number.json (100%)
 rename examples/json/{extractor => lockfile}/value_object.json (100%)
 rename examples/json/{extractor => lockfile}/value_string.json (100%)

diff --git a/circuits/test/http/codegen.test.ts b/circuits/test/http/codegen.test.ts
index 580f72e..496176a 100644
--- a/circuits/test/http/codegen.test.ts
+++ b/circuits/test/http/codegen.test.ts
@@ -75,10 +75,10 @@ describe("HTTP :: Codegen :: Request", async () => {
         // generate extractor circuit using codegen
         await executeCodegen(`${lockfile}.json`, lockfile);
 
-        const lockData = await readLockFile<Request>(`${lockfile}.json`);
+        const lockData = readLockFile<Request>(`${lockfile}.json`);
         console.log("lockData: ", JSON.stringify(lockData));
 
-        const input = await readHTTPInputFile(`${inputfile}`).input;
+        const input = readHTTPInputFile(`${inputfile}`).input;
 
         const headers = getHeaders(lockData);
         const params = [input.length, lockData.method.length, lockData.target.length, lockData.version.length];
@@ -117,9 +117,9 @@ describe("HTTP :: Codegen :: Request", async () => {
         // generate extractor circuit using codegen
         await executeCodegen(`${lockfile}.json`, lockfile);
 
-        const lockData = await readLockFile<Request>(`${lockfile}.json`);
+        const lockData = readLockFile<Request>(`${lockfile}.json`);
 
-        const input = await readHTTPInputFile(`${inputfile}`).input
+        const input = readHTTPInputFile(`${inputfile}`).input
 
         const headers = getHeaders(lockData);
         const params = [input.length, lockData.method.length, lockData.target.length, lockData.version.length];
@@ -163,10 +163,10 @@ describe("HTTP :: Codegen :: Response", async () => {
         // generate extractor circuit using codegen
         await executeCodegen(`${lockfile}.json`, lockfile);
 
-        const lockData = await readLockFile<Response>(`${lockfile}.json`);
+        const lockData = readLockFile<Response>(`${lockfile}.json`);
         console.log("lockData: ", JSON.stringify(lockData));
 
-        const http = await readHTTPInputFile(`${inputfile}`);
+        const http = readHTTPInputFile(`${inputfile}`);
         const input = http.input;
 
         const headers = getHeaders(lockData);
@@ -209,9 +209,9 @@ describe("HTTP :: Codegen :: Response", async () => {
         // generate extractor circuit using codegen
         await executeCodegen(`${lockfile}.json`, lockfile);
 
-        const lockData = await readLockFile<Response>(`${lockfile}.json`);
+        const lockData = readLockFile<Response>(`${lockfile}.json`);
 
-        const http = await readHTTPInputFile(`${inputfile}`);
+        const http = readHTTPInputFile(`${inputfile}`);
         const input = http.input;
 
         const headers = getHeaders(lockData);
diff --git a/docs/http.md b/docs/http.md
new file mode 100644
index 0000000..e69de29
diff --git a/docs/pabuild.md b/docs/pabuild.md
index c218d3b..97191c2 100644
--- a/docs/pabuild.md
+++ b/docs/pabuild.md
@@ -1,15 +1,15 @@
-# `wpbuild` CLI Tool
-This repository contains a small Rust CLI tool called `wpbuild`. 
+# `pabuild` CLI Tool
+This repository contains a small Rust CLI tool called `pabuild`.
 
-## Install `wpbuild`
+## Install `pabuild`
 From the root of this repository, run:
 ```sh
 cargo install --path .
 ```
-to install the `wpbuild` binary.
+to install the `pabuild` binary.
 You can see a help menu with the subcommands by:
 ```sh
-wpbuild --help
+pabuild --help
 ```
 
 ## Witnessgen
@@ -21,19 +21,19 @@ It can process and generate JSON files to be used for these circuits.
 
 ### Examples
 **JSON Parsing:**
-If we have a given JSON file we want to parse such as `examples/json/test/example.json` for the `json-parser` circuit (see `circuits.json`), then we can:
+If we have a given JSON file we want to parse such as [`examples/json/test/example.json`](../examples/json/test/example.json) for the `json-parser` circuit (see [`circuits.json`](../circuits.json)), then we can:
 
 ```sh
-pabuild witness json --input-file examples/json/test/example.json --output-dir inputs/json-parser --output-filename input.json
+pabuild witness json --input-file examples/json/test/example.json --output-dir inputs/json-parser --output-filename input.json json
 ```
 
 Afterwards, you can run `npx circomkit compile json-parser` then `circomkit witness json-parser input`.
 
 **HTTP Parsing:**
-If we have a given HTTP request/response (as a file) we want to parse such as `examples/http/get_request.http` for the `http-parser` circuit (see `circuits.json`), then we can:
+If we have a given HTTP request/response (as a file) we want to parse such as [`examples/http/get_request.http`](../examples/http/get_request.http) for the `http-parser` circuit (see `circuits.json`), then we can:
 
 ```sh
-pabuild witness http --input-file examples/json/get_request.http --output-dir inputs/http-parser --output-filename input.json
+pabuild witness http --input-file examples/json/get_request.http --output-dir inputs/http-parser --output-filename input.json http
 ```
 
 Afterwards, you can run `npx circomkit compile http-parser` then `circomkit witness http-parser input`.
@@ -44,25 +44,31 @@ Afterwards, you can run `npx circomkit compile http-parser` then `circomkit witn
 JSON extractor circuit is generated using rust to handle arbitrary keys and array indices.
 
 Run:
-```bash
-cargo run --bin codegen -- --help
+```sh
+pabuild json --help
 ```
 to get options:
 ```
-Usage: codegen [OPTIONS] --json-file <JSON_FILE>
+Usage: pabuild json [OPTIONS] --template <TEMPLATE>
 
 Options:
-  -j, --json-file <JSON_FILE>              Path to the JSON file
+  -t, --template <TEMPLATE>                Path to the JSON file selective-disclosure template
   -o, --output-filename <OUTPUT_FILENAME>  Output circuit file name [default: extractor]
+  -d, --debug                              Optional circuit debug logs
+  -h, --help                               Print help
 ```
-Takes input 2 arguments:
-- `json-file`: input json file. Examples are located in [codegen](./examples/json/test/codegen/)
-- `output-filename`: circuit filename to save. Located in [circuits/main](./circuits/main/). If not given, defaults to `extractor.circom`.
+Takes 3 input arguments:
+- `template`: input json file. Examples are located in [extractor](../examples/extractor/).
+  - Should contain only two keys:
+    - `keys`: list of all the keys in the input json
+    - `value_type`: Currently only two value types are supported: `String`,`Number`.
+- `output-filename`: circuit filename to save. Located in [circuits/main](../circuits/main/). If not given, defaults to `extractor.circom`.
+- `debug`: Optional debug logs for parser and extractor output.
 
 To test an end-to-end JSON extraction proof:
 - Run codegen to generate circuits. Replace `value_string` with input filename.
-   ```bash
-   cargo run --bin codegen -- --json-file ./examples/json/test/codegen/value_string.json --output-filename value_string
+   ```sh
+   pabuild json --template examples/json/extractor/value_string.json --output-filename value_string
    ```
 
 - Compile circom circuit using
@@ -70,7 +76,7 @@ To test an end-to-end JSON extraction proof:
    circom ./circuits/main/value_string.circom --r1cs --wasm
    ```
 
-- To use circomkit: add circuit config to [circuits.json](./circuits.json). and input file to [inputs](./inputs/)
+- To use circomkit: add circuit config to [circuits.json](../circuits.json). and input file to [inputs](../inputs/)
 
 - Generate witness:
    ```bash
diff --git a/examples/json/extractor/two_keys.json b/examples/json/lockfile/two_keys.json
similarity index 100%
rename from examples/json/extractor/two_keys.json
rename to examples/json/lockfile/two_keys.json
diff --git a/examples/json/extractor/value_array_nested.json b/examples/json/lockfile/value_array_nested.json
similarity index 100%
rename from examples/json/extractor/value_array_nested.json
rename to examples/json/lockfile/value_array_nested.json
diff --git a/examples/json/extractor/value_array_number.json b/examples/json/lockfile/value_array_number.json
similarity index 100%
rename from examples/json/extractor/value_array_number.json
rename to examples/json/lockfile/value_array_number.json
diff --git a/examples/json/extractor/value_array_object.json b/examples/json/lockfile/value_array_object.json
similarity index 100%
rename from examples/json/extractor/value_array_object.json
rename to examples/json/lockfile/value_array_object.json
diff --git a/examples/json/extractor/value_array_string.json b/examples/json/lockfile/value_array_string.json
similarity index 100%
rename from examples/json/extractor/value_array_string.json
rename to examples/json/lockfile/value_array_string.json
diff --git a/examples/json/extractor/value_number.json b/examples/json/lockfile/value_number.json
similarity index 100%
rename from examples/json/extractor/value_number.json
rename to examples/json/lockfile/value_number.json
diff --git a/examples/json/extractor/value_object.json b/examples/json/lockfile/value_object.json
similarity index 100%
rename from examples/json/extractor/value_object.json
rename to examples/json/lockfile/value_object.json
diff --git a/examples/json/extractor/value_string.json b/examples/json/lockfile/value_string.json
similarity index 100%
rename from examples/json/extractor/value_string.json
rename to examples/json/lockfile/value_string.json
diff --git a/src/json.rs b/src/json.rs
index 9d6982c..7a21e31 100644
--- a/src/json.rs
+++ b/src/json.rs
@@ -3,6 +3,8 @@ use std::{
     str::FromStr,
 };
 
+use std::collections::HashMap;
+
 use super::*;
 
 const PRAGMA: &str = "pragma circom 2.1.9;\n\n";
@@ -23,12 +25,25 @@ pub enum Key {
 }
 
 #[derive(Debug, Deserialize)]
-pub struct Data {
+pub struct JsonLockfile {
     keys: Vec<Key>,
     value_type: ValueType,
 }
 
-fn extract_string(data: Data, circuit_buffer: &mut String, debug: bool) {
+impl JsonLockfile {
+    pub fn as_bytes(&self) -> HashMap<String, Vec<u8>> {
+        let mut keys = HashMap::<String, Vec<u8>>::new();
+        for (i, key) in self.keys.iter().enumerate() {
+            if let Key::String(key) = key {
+                let key_name = format!("key{}", i + 1);
+                keys.insert(key_name, key.as_bytes().to_vec());
+            }
+        }
+        keys
+    }
+}
+
+fn extract_string(data: JsonLockfile, circuit_buffer: &mut String, debug: bool) {
     *circuit_buffer += "template ExtractStringValue(DATA_BYTES, MAX_STACK_HEIGHT, ";
     for (i, key) in data.keys.iter().enumerate() {
         match key {
@@ -93,7 +108,7 @@ fn extract_string(data: Data, circuit_buffer: &mut String, debug: bool) {
 "#;
 }
 
-fn extract_number(data: Data, circuit_buffer: &mut String, debug: bool) {
+fn extract_number(data: JsonLockfile, circuit_buffer: &mut String, debug: bool) {
     *circuit_buffer += "template ExtractNumValue(DATA_BYTES, MAX_STACK_HEIGHT, ";
     for (i, key) in data.keys.iter().enumerate() {
         match key {
@@ -169,7 +184,7 @@ fn extract_number(data: Data, circuit_buffer: &mut String, debug: bool) {
 }
 
 fn build_json_circuit(
-    data: Data,
+    data: JsonLockfile,
     output_filename: String,
     debug: bool,
 ) -> Result<(), Box<dyn std::error::Error>> {
@@ -523,7 +538,7 @@ fn build_json_circuit(
 
 pub fn json_circuit(args: JsonArgs) -> Result<(), Box<dyn std::error::Error>> {
     let data = std::fs::read(&args.template)?;
-    let json_data: Data = serde_json::from_slice(&data)?;
+    let json_data: JsonLockfile = serde_json::from_slice(&data)?;
 
     build_json_circuit(json_data, args.output_filename, args.debug)?;
 
diff --git a/src/main.rs b/src/main.rs
index e363a9b..5a1757d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,13 +15,14 @@ pub struct Args {
 
 #[derive(Subcommand, Debug)]
 pub enum Command {
-    Witness(WitnessArgs),
+    ParserWitness(ParserWitnessArgs),
+    ExtractorWitness(ExtractorWitnessArgs),
     Json(JsonArgs),
     Http(HttpArgs),
 }
 
 #[derive(Parser, Debug)]
-pub struct WitnessArgs {
+pub struct ParserWitnessArgs {
     #[arg(global = true, value_enum)]
     subcommand: WitnessSubcommand,
 
@@ -38,6 +39,28 @@ pub struct WitnessArgs {
     output_filename: String,
 }
 
+#[derive(Parser, Debug)]
+pub struct ExtractorWitnessArgs {
+    #[arg(global = true, value_enum)]
+    subcommand: WitnessSubcommand,
+
+    /// Path to the JSON file
+    #[arg(global = true, long)]
+    input_file: PathBuf,
+
+    /// Path to the lockfile
+    #[arg(global = true, long)]
+    lockfile: PathBuf,
+
+    /// Output directory (will be created if it doesn't exist)
+    #[arg(global = true, long, default_value = ".")]
+    output_dir: PathBuf,
+
+    /// Output filename (will be created if it doesn't exist)
+    #[arg(global = true, long, default_value = "output.json")]
+    output_filename: String,
+}
+
 #[derive(clap::ValueEnum, Clone, Debug)]
 pub enum WitnessSubcommand {
     Json,
@@ -76,8 +99,9 @@ pub struct HttpArgs {
 
 pub fn main() -> Result<(), Box<dyn Error>> {
     match Args::parse().command {
-        Command::Witness(args) => witness::witness(args),
+        Command::ParserWitness(args) => witness::parser_witness(args),
         Command::Json(args) => json::json_circuit(args),
         Command::Http(args) => http::http_circuit(args),
+        Command::ExtractorWitness(args) => witness::extractor_witness(args),
     }
 }
diff --git a/src/witness.rs b/src/witness.rs
index 346e4db..522e6c8 100644
--- a/src/witness.rs
+++ b/src/witness.rs
@@ -1,17 +1,42 @@
+use json::JsonLockfile;
+
 use super::*;
-use std::io::Write;
+use std::{collections::HashMap, io::Write};
 
 #[derive(serde::Serialize)]
 pub struct Witness {
     data: Vec<u8>,
 }
 
-pub fn witness(args: WitnessArgs) -> Result<(), Box<dyn std::error::Error>> {
+#[derive(serde::Serialize)]
+pub struct ExtractorWitness {
+    data: Vec<u8>,
+    keys: HashMap<String, Vec<u8>>,
+}
+
+fn print_boxed_output(lines: Vec<String>) {
+    // Determine the maximum length of the lines
+    let max_length = lines.iter().map(|line| line.len()).max().unwrap_or(0);
+
+    // Characters for the box
+    let top_border = format!("┌{}┐", "─".repeat(max_length + 2));
+    let bottom_border = format!("└{}┘", "─".repeat(max_length + 2));
+
+    // Print the box with content
+    println!("{}", top_border);
+    for line in lines {
+        println!("│ {:<width$} │", line, width = max_length);
+    }
+    println!("{}", bottom_border);
+}
+
+pub fn parser_witness(args: ParserWitnessArgs) -> Result<(), Box<dyn std::error::Error>> {
     let data = match &args.subcommand {
         WitnessSubcommand::Json => std::fs::read(args.input_file)?,
         WitnessSubcommand::Http => {
             let mut data = std::fs::read(args.input_file)?;
             let mut i = 0;
+            // convert LF to CRLF
             while i < data.len() {
                 if data[i] == 10 && (i == 0 || data[i - 1] != 13) {
                     data.insert(i, 13);
@@ -44,18 +69,53 @@ pub fn witness(args: WitnessArgs) -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
-fn print_boxed_output(lines: Vec<String>) {
-    // Determine the maximum length of the lines
-    let max_length = lines.iter().map(|line| line.len()).max().unwrap_or(0);
+fn read_input_file_as_bytes(
+    file_type: WitnessSubcommand,
+    file_path: PathBuf,
+) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+    match file_type {
+        WitnessSubcommand::Json => Ok(std::fs::read(file_path)?),
+        WitnessSubcommand::Http => {
+            let mut data = std::fs::read(file_path)?;
+            let mut i = 0;
+            // convert LF to CRLF
+            while i < data.len() {
+                if data[i] == 10 && (i == 0 || data[i - 1] != 13) {
+                    data.insert(i, 13);
+                    i += 2;
+                } else {
+                    i += 1;
+                }
+            }
+            Ok(data)
+        }
+    }
+}
+pub fn extractor_witness(args: ExtractorWitnessArgs) -> Result<(), Box<dyn std::error::Error>> {
+    let input_data = read_input_file_as_bytes(args.subcommand, args.input_file)?;
 
-    // Characters for the box
-    let top_border = format!("┌{}┐", "─".repeat(max_length + 2));
-    let bottom_border = format!("└{}┘", "─".repeat(max_length + 2));
+    let lockfile_data = std::fs::read(&args.lockfile)?;
+    let lockfile: JsonLockfile = serde_json::from_slice(&lockfile_data)?;
 
-    // Print the box with content
-    println!("{}", top_border);
-    for line in lines {
-        println!("│ {:<width$} │", line, width = max_length);
+    let witness = ExtractorWitness {
+        data: input_data.clone(),
+        keys: lockfile.as_bytes(),
+    };
+
+    if !args.output_dir.exists() {
+        std::fs::create_dir_all(&args.output_dir)?;
     }
-    println!("{}", bottom_border);
+
+    let output_file = args.output_dir.join(args.output_filename);
+    let mut file = std::fs::File::create(output_file)?;
+    file.write_all(serde_json::to_string_pretty(&witness)?.as_bytes())?;
+
+    // Prepare lines to print
+    let mut lines = Vec::new();
+    lines.push(format!("Data length: {}", input_data.len()));
+
+    // Print the output inside a nicely formatted box
+    print_boxed_output(lines);
+
+    Ok(())
 }