Contractless/src/bin/recreate_wallet_from_image.rs

use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use blockchain::common::cli_prompts::prompt_hidden_nonempty;
use blockchain::common::network_startup::get_connections;
use blockchain::decode_image_and_extract_text;
use blockchain::decrypts;
use blockchain::env;
use blockchain::fs;
use blockchain::read_to_string;
use blockchain::records::memory::response_channels::generate_uid;
use blockchain::standalone_tools::connections::handshake;
use blockchain::stdout;
use blockchain::tilde;
use blockchain::wallets::structures::{SavedWallet, Wallet};
use blockchain::AsyncWriteExt;
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;
use std::path::PathBuf;

fn display_vanity_address(short_address: &str) -> Option<String> {
    let (payload, network_suffix) = short_address.rsplit_once('.')?;
    let trimmed_payload = payload.trim_start();
    if trimmed_payload.is_empty() {
        return None;
    }
    Some(format!("{trimmed_payload}.{network_suffix}"))
}

async fn lookup_registered_vanity_address(
    short_address: &str,
    long_address: &str,
    private_key: &str,
) -> Option<String> {
    let hashmap_key = generate_uid();
    let rpc_command = 40;
    let connections = get_connections().await;

    for conn in connections {
        let Ok(socket_address) = conn.parse() else {
            continue;
        };

        let result = handshake::connect_and_handshake(
            socket_address,
            short_address.to_string(),
            rpc_command,
            handshake::HandshakeWallet::WalletParts {
                long_address: long_address.to_string(),
                private_key: private_key.to_string(),
            },
            hashmap_key,
        )
        .await;

        if let Ok(response) = result {
            if response.is_empty() {
                return None;
            }

            let response_text = String::from_utf8_lossy(&response);
            let trimmed = response_text.trim();
            if trimmed.is_empty() || trimmed.starts_with("error:") {
                return None;
            }

            return display_vanity_address(trimmed);
        }
    }

    None
}

pub fn load_wallet(private_key_image: String, wallet_key: String) -> String {
    let private_key: String;
    if let Some(encrypted_text) = decode_image_and_extract_text(&private_key_image) {
        // Decrypt the encrypted text to get the real private key
        if let Some(decrypted_private_key) = decrypts(&encrypted_text, Some(&wallet_key)) {
            // Update the wallet's private key with the decrypted private 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
}

#[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}")),
    }
}

fn prompt_for_filename(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) => {
            let trimmed = line.trim();
            if trimmed.is_empty() {
                Err("Filename cannot be empty".to_string())
            } else {
                Ok(trimmed.to_string())
            }
        }
        Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
            Err("Input cancelled".to_string())
        }
        Err(err) => Err(format!("Failed to read filename: {err}")),
    }
}

async fn load_private_key_image_input(path: &str) -> Result<String, String> {
    let expanded_path = tilde(path).to_string();

    if expanded_path.to_ascii_lowercase().ends_with(".png") {
        let image_bytes =
            fs::read(&expanded_path).map_err(|e| format!("Failed to read PNG file: {e}"))?;
        return Ok(STANDARD.encode(image_bytes));
    }

    let file_contents = read_to_string(&expanded_path)
        .await
        .map_err(|e| format!("Failed to read Base64 string file: {e}"))?;

    Ok(file_contents.trim().to_string())
}

#[tokio::main]
async fn main() {
    let mut stdout = stdout(); // Get Tokio's stdout

    // Collect command-line arguments
    let args: Vec<String> = env::args().collect();

    let base64_file_path: String;
    let output_path: String;

    if args.len() > 1 && args.len() != 3 && args.len() != 4 {
        println!(
            "Usage: ./recreate_wallet_from_image <base64_or_png_file_path> <output_wallet_file>"
        );
        println!("   or: ./recreate_wallet_from_image <base64_or_png_file_path> <output_wallet_dir> <output_wallet_filename>");
        return;
    }

    if args.len() == 4 {
        base64_file_path = args[1].clone();
        output_path = PathBuf::from(&args[2])
            .join(&args[3])
            .to_string_lossy()
            .to_string();
    } else if args.len() == 3 {
        base64_file_path = args[1].clone();
        output_path = args[2].clone();
    } else {
        stdout.flush().await.expect("Failed to flush stdout");
        base64_file_path =
            match prompt_for_path("Please enter the path to the Base64 or PNG image file: ") {
                Ok(path) => path,
                Err(err) => {
                    eprintln!("{err}");
                    return;
                }
            };

        let output_dir_path = match prompt_for_path(
            "Please enter the directory path where the rebuilt wallet JSON should be written: ",
        ) {
            Ok(path) => path,
            Err(err) => {
                eprintln!("{err}");
                return;
            }
        };

        let output_filename =
            match prompt_for_filename("Please enter the filename for the rebuilt wallet JSON: ") {
                Ok(filename) => filename,
                Err(err) => {
                    eprintln!("{err}");
                    return;
                }
            };

        let output_dir = tilde(&output_dir_path).to_string();
        output_path = PathBuf::from(output_dir)
            .join(output_filename)
            .to_string_lossy()
            .to_string();
    }

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

    let base64_string = match load_private_key_image_input(&base64_file_path).await {
        Ok(contents) => contents,
        Err(err) => {
            eprintln!("{err}");
            return;
        }
    };
    let expanded_output_path = tilde(&output_path).to_string();

    // Generate private key from Base64 string
    let private_key = load_wallet(base64_string.clone(), wallet_key);

    match Wallet::regenerate_public_key(&private_key) {
        Ok(public_key_bytes) => {
            let public_key = blockchain::encode(public_key_bytes.clone());
            let long_address = Wallet::generate_address(&public_key);
            let long_address_bytes = Wallet::long_address_to_bytes(long_address.clone());
            let short_address_bytes =
                Wallet::long_address_bytes_to_short_address_bytes(&long_address_bytes)
                    .expect("Failed to derive short address bytes");
            let short_address = Wallet::bytes_to_short_address(&short_address_bytes)
                .expect("Failed to encode short address");
            let vanity_address =
                lookup_registered_vanity_address(&short_address, &long_address, &private_key).await;
            let saved_wallet = SavedWallet {
                long_address,
                short_address,
                vanity_address,
                public_key,
                private_key: base64_string,
            };
            let output_str = serde_json::to_string_pretty(&saved_wallet).unwrap();
            fs::write(&expanded_output_path, output_str)
                .expect("Failed to write rebuilt wallet file");
            println!("Wallet written to {expanded_output_path}");
        }
        Err(e) => println!("Error: {e}"),
    }
}