Contractless/src/bin/private_key_from_image.rs

use blockchain::common::cli_prompts::prompt_hidden_nonempty;
use blockchain::decode_image_and_extract_text;
use blockchain::decrypts;
use blockchain::env;
use blockchain::from_str;
use blockchain::fs;
use blockchain::tilde;
use blockchain::Value;
use rustyline::completion::FilenameCompleter;
use rustyline::error::ReadlineError;
use rustyline::{history::DefaultHistory, CompletionType, Config, Editor};
use rustyline_derive::Completer;
use rustyline_derive::Helper as RustyHelper;
use rustyline_derive::Highlighter as RustyHighlighter;
use rustyline_derive::Hinter as RustyHinter;
use rustyline_derive::Validator as RustyValidator;

pub async fn decode_private_key(base64_image: String, wallet_key: String) -> String {
    let private_key: String;
    if let Some(encrypted_text) = decode_image_and_extract_text(&base64_image) {
        // Decrypt the encrypted text to get the real private key
        if let Some(decrypted_private_key) = decrypts(&encrypted_text, Some(&wallet_key)) {
            private_key = decrypted_private_key;
        } else {
            eprintln!("Decryption of private key failed.");
            panic!();
        }
    } else {
        eprintln!("Failed to decode the image and extract text.");
        panic!();
    }
    private_key
}

fn extract_private_key(contents: &str) -> Result<String, String> {
    let trimmed = contents.trim();
    if trimmed.is_empty() {
        return Err("Wallet/image file is empty".to_string());
    }

    if trimmed.starts_with('{') {
        let value: Value =
            from_str(trimmed).map_err(|e| format!("Failed to parse wallet JSON: {e}"))?;
        let private_key = value
            .get("private_key")
            .or_else(|| value.get("privkey"))
            .and_then(|v| v.as_str())
            .ok_or_else(|| {
                "Wallet JSON does not contain a \"private_key\" or \"privkey\" field".to_string()
            })?;
        return Ok(private_key.trim().to_string());
    }

    Ok(trimmed.to_string())
}

#[derive(RustyHelper, Completer, RustyHinter, RustyHighlighter, RustyValidator)]
struct PathHelper {
    #[rustyline(Completer)]
    completer: FilenameCompleter,
}

fn prompt_for_path(prompt: &str) -> Result<String, String> {
    let config = Config::builder()
        .completion_type(CompletionType::List)
        .build();
    let mut editor =
        Editor::<PathHelper, DefaultHistory>::with_config(config).map_err(|e| e.to_string())?;
    editor.set_helper(Some(PathHelper {
        completer: FilenameCompleter::new(),
    }));

    match editor.readline(prompt) {
        Ok(line) => Ok(line.trim().to_string()),
        Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
            Err("Input cancelled".to_string())
        }
        Err(err) => Err(format!("Failed to read file path: {err}")),
    }
}

#[tokio::main]
async fn main() {
    // Collect command-line arguments
    let args: Vec<String> = env::args().collect();

    let base64_file_path = if args.len() > 1 {
        args[1].clone()
    } else {
        // Prompt the user for file path and encryption key
        match prompt_for_path(
            "Please enter the path to the file containing the wallet/image data: ",
        ) {
            Ok(path) => path,
            Err(err) => {
                eprintln!("{err}");
                return;
            }
        }
    };

    let wallet_key = prompt_hidden_nonempty(
        "Please enter your encryption key: ",
        "Wallet key cannot be empty. Please try again.",
    )
    .await;

    // Read the Base64 string from the specified file
    let expanded_path = tilde(&base64_file_path).to_string();
    let file_contents = fs::read_to_string(&expanded_path)
        .expect("Failed to read Base64 string file")
        .to_string();

    let base64_string = match extract_private_key(&file_contents) {
        Ok(value) => value,
        Err(err) => {
            eprintln!("{err}");
            return;
        }
    };

    // Decode the private key from Base64 string
    let private_key = decode_private_key(base64_string, wallet_key).await;

    // Print the private key in text format
    println!("{private_key}");
}