darkfi/bin/drk2/src/main.rs

/* This file is part of DarkFi (https://dark.fi)
 *
 * Copyright (C) 2020-2024 Dyne.org foundation
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use std::{
    fs,
    io::{stdin, Read},
    process::exit,
    str::FromStr,
    sync::Arc,
    time::Instant,
};

use prettytable::{format, row, Table};
use smol::stream::StreamExt;
use structopt_toml::{serde::Deserialize, structopt::StructOpt, StructOptToml};
use url::Url;

use darkfi::{
    async_daemonize, cli_desc,
    rpc::{client::RpcClient, jsonrpc::JsonRequest, util::JsonValue},
    tx::Transaction,
    util::{parse::encode_base10, path::expand_path},
    Result,
};
use darkfi_money_contract::model::Coin;
use darkfi_sdk::{
    crypto::TokenId,
    pasta::{group::ff::PrimeField, pallas},
};
use darkfi_serial::{deserialize, serialize};

/// Error codes
mod error;

/// darkfid JSON-RPC related methods
mod rpc;

/// CLI utility functions
mod cli_util;
use cli_util::kaching;

/// Wallet functionality related to Money
mod money;
use money::BALANCE_BASE10_DECIMALS;

/// Wallet functionality related to Dao
mod dao;

/// Wallet functionality related to transactions history
mod txs_history;

/// Wallet database operations handler
mod walletdb;
use walletdb::{WalletDb, WalletPtr};

const CONFIG_FILE: &str = "drk_config.toml";
const CONFIG_FILE_CONTENTS: &str = include_str!("../drk_config.toml");

#[derive(Clone, Debug, Deserialize, StructOpt, StructOptToml)]
#[serde(default)]
#[structopt(name = "drk", about = cli_desc!())]
struct Args {
    #[structopt(short, long)]
    /// Configuration file to use
    config: Option<String>,

    #[structopt(long, default_value = "~/.local/darkfi/drk/wallet.db")]
    /// Path to wallet database
    wallet_path: String,

    #[structopt(long, default_value = "changeme")]
    /// Password for the wallet database
    wallet_pass: String,

    #[structopt(short, long, default_value = "tcp://127.0.0.1:8340")]
    /// darkfid JSON-RPC endpoint
    endpoint: Url,

    #[structopt(subcommand)]
    /// Sub command to execute
    command: Subcmd,

    #[structopt(short, long)]
    /// Set log file to ouput into
    log: Option<String>,

    #[structopt(short, parse(from_occurrences))]
    /// Increase verbosity (-vvv supported)
    verbose: u8,
}

#[derive(Clone, Debug, Deserialize, StructOpt)]
enum Subcmd {
    /// Fun
    Kaching,

    /// Send a ping request to the darkfid RPC endpoint
    Ping,

    // TODO: shell completions
    /// Wallet operations
    Wallet {
        #[structopt(long)]
        /// Initialize wallet database
        initialize: bool,

        #[structopt(long)]
        /// Generate a new keypair in the wallet
        keygen: bool,

        #[structopt(long)]
        /// Query the wallet for known balances
        balance: bool,

        #[structopt(long)]
        /// Get the default address in the wallet
        address: bool,

        #[structopt(long)]
        /// Print all the addresses in the wallet
        addresses: bool,

        #[structopt(long)]
        /// Set the default address in the wallet
        default_address: Option<usize>,

        #[structopt(long)]
        /// Print all the secret keys from the wallet
        secrets: bool,

        #[structopt(long)]
        /// Import secret keys from stdin into the wallet, separated by newlines
        import_secrets: bool,

        #[structopt(long)]
        /// Print the Merkle tree in the wallet
        tree: bool,

        #[structopt(long)]
        /// Print all the coins in the wallet
        coins: bool,
    },

    /// Unspend a coin
    Unspend {
        /// base58-encoded coin to mark as unspent
        coin: String,
    },

    // TODO: Transfer

    // TODO: OTC
    /// Inspect a transaction from stdin
    Inspect,

    /// Read a transaction from stdin and broadcast it
    Broadcast,

    /// This subscription will listen for incoming blocks from darkfid and look
    /// through their transactions to see if there's any that interest us.
    /// With `drk` we look at transactions calling the money contract so we can
    /// find coins sent to us and fill our wallet with the necessary metadata.
    Subscribe,

    // TODO: DAO
    /// Scan the blockchain and parse relevant transactions
    Scan {
        #[structopt(long)]
        /// Reset Merkle tree and start scanning from first block
        reset: bool,

        #[structopt(long)]
        /// List all available checkpoints
        list: bool,

        #[structopt(long)]
        /// Reset Merkle tree to checkpoint index and start scanning
        checkpoint: Option<u64>,
    },

    /// Explorer related subcommands
    Explorer {
        #[structopt(subcommand)]
        /// Sub command to execute
        command: ExplorerSubcmd,
    },

    /// Manage Token aliases
    Alias {
        #[structopt(subcommand)]
        /// Sub command to execute
        command: AliasSubcmd,
    },
    // TODO: Token
}

#[derive(Clone, Debug, Deserialize, StructOpt)]
enum ExplorerSubcmd {
    /// Fetch a blockchain transaction by hash
    FetchTx {
        /// Transaction hash
        tx_hash: String,

        #[structopt(long)]
        /// Print the full transaction information
        full: bool,

        #[structopt(long)]
        /// Encode transaction to base58
        encode: bool,
    },

    /// Read a transaction from stdin and simulate it
    SimulateTx,

    /// Fetch broadcasted transactions history
    TxsHistory {
        /// Fetch specific history record (optional)
        tx_hash: Option<String>,

        #[structopt(long)]
        /// Encode specific history record transaction
        /// to base58.
        encode: bool,
    },
}

#[derive(Clone, Debug, Deserialize, StructOpt)]
enum AliasSubcmd {
    /// Create a Token alias
    Add {
        /// Token alias
        alias: String,

        /// Token to create alias for
        token: String,
    },

    /// Print alias info of optional arguments.
    /// If no argument is provided, list all the aliases in the wallet.
    Show {
        /// Token alias to search for
        #[structopt(short, long)]
        alias: Option<String>,

        /// Token to search alias for
        #[structopt(short, long)]
        token: Option<String>,
    },

    /// Remove a Token alias
    Remove {
        /// Token alias to remove
        alias: String,
    },
}

/// CLI-util structure
pub struct Drk {
    /// Wallet database operations handler
    pub wallet: WalletPtr,
    /// JSON-RPC client to execute requests to darkfid daemon
    pub rpc_client: RpcClient,
}

impl Drk {
    async fn new(
        wallet_path: String,
        wallet_pass: String,
        endpoint: Url,
        ex: Arc<smol::Executor<'static>>,
    ) -> Result<Self> {
        // Initialize wallet
        let wallet_path = expand_path(&wallet_path)?;
        if !wallet_path.exists() {
            if let Some(parent) = wallet_path.parent() {
                fs::create_dir_all(parent)?;
            }
        }
        let wallet = match WalletDb::new(Some(wallet_path), Some(&wallet_pass)) {
            Ok(w) => w,
            Err(e) => {
                eprintln!("Error initializing wallet: {e:?}");
                exit(2);
            }
        };

        // Initialize rpc client
        let rpc_client = RpcClient::new(endpoint, ex).await?;

        Ok(Self { wallet, rpc_client })
    }

    /// Initialize wallet with tables for drk
    async fn initialize_wallet(&self) -> Result<()> {
        let wallet_schema = include_str!("../wallet.sql");
        if let Err(e) = self.wallet.exec_batch_sql(wallet_schema).await {
            eprintln!("Error initializing wallet: {e:?}");
            exit(2);
        }

        Ok(())
    }

    /// Auxilliary function to ping configured darkfid daemon for liveness.
    async fn ping(&self) -> Result<()> {
        eprintln!("Executing ping request to darkfid...");
        let latency = Instant::now();
        let req = JsonRequest::new("ping", JsonValue::Array(vec![]));
        let rep = self.rpc_client.oneshot_request(req).await?;
        let latency = latency.elapsed();
        eprintln!("Got reply: {rep:?}");
        eprintln!("Latency: {latency:?}");
        Ok(())
    }
}

async_daemonize!(realmain);
async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
    match args.command {
        Subcmd::Kaching => {
            kaching().await;
            Ok(())
        }

        Subcmd::Ping => {
            let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;
            drk.ping().await
        }

        Subcmd::Wallet {
            initialize,
            keygen,
            balance,
            address,
            addresses,
            default_address,
            secrets,
            import_secrets,
            tree,
            coins,
        } => {
            if !initialize &&
                !keygen &&
                !balance &&
                !address &&
                !addresses &&
                default_address.is_none() &&
                !secrets &&
                !tree &&
                !coins &&
                !import_secrets
            {
                eprintln!("Error: You must use at least one flag for this subcommand");
                eprintln!("Run with \"wallet -h\" to see the subcommand usage.");
                exit(2);
            }

            let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;

            if initialize {
                drk.initialize_wallet().await?;
                if let Err(e) = drk.initialize_money().await {
                    eprintln!("Failed to initialize Money: {e:?}");
                    exit(2);
                }
                if let Err(e) = drk.initialize_dao().await {
                    eprintln!("Failed to initialize DAO: {e:?}");
                    exit(2);
                }
                return Ok(())
            }

            if keygen {
                if let Err(e) = drk.money_keygen().await {
                    eprintln!("Failed to generate keypair: {e:?}");
                    exit(2);
                }
                return Ok(())
            }

            if balance {
                let balmap = drk.money_balance().await?;

                let aliases_map = drk.get_aliases_mapped_by_token().await?;

                // Create a prettytable with the new data:
                let mut table = Table::new();
                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
                table.set_titles(row!["Token ID", "Aliases", "Balance"]);
                for (token_id, balance) in balmap.iter() {
                    let aliases = match aliases_map.get(token_id) {
                        Some(a) => a,
                        None => "-",
                    };

                    table.add_row(row![
                        token_id,
                        aliases,
                        encode_base10(*balance, BALANCE_BASE10_DECIMALS)
                    ]);
                }

                if table.is_empty() {
                    eprintln!("No unspent balances found");
                } else {
                    eprintln!("{table}");
                }

                return Ok(())
            }

            if address {
                let address = match drk.default_address().await {
                    Ok(a) => a,
                    Err(e) => {
                        eprintln!("Failed to fetch default address: {e:?}");
                        exit(2);
                    }
                };

                eprintln!("{address}");

                return Ok(())
            }

            if addresses {
                let addresses = drk.addresses().await?;

                // Create a prettytable with the new data:
                let mut table = Table::new();
                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
                table.set_titles(row!["Key ID", "Public Key", "Secret Key", "Is Default"]);
                for (key_id, public_key, secret_key, is_default) in addresses {
                    let is_default = match is_default {
                        1 => "*",
                        _ => "",
                    };
                    table.add_row(row![key_id, public_key, secret_key, is_default]);
                }

                if table.is_empty() {
                    eprintln!("No addresses found");
                } else {
                    eprintln!("{table}");
                }

                return Ok(())
            }

            if let Some(idx) = default_address {
                if let Err(e) = drk.set_default_address(idx).await {
                    eprintln!("Failed to set default address: {e:?}");
                    exit(2);
                }
                return Ok(())
            }

            if secrets {
                let v = drk.get_money_secrets().await?;

                for i in v {
                    eprintln!("{i}");
                }

                return Ok(())
            }

            if import_secrets {
                let mut secrets = vec![];
                let lines = stdin().lines();
                for (i, line) in lines.enumerate() {
                    if let Ok(line) = line {
                        let bytes = bs58::decode(&line.trim()).into_vec()?;
                        let Ok(secret) = deserialize(&bytes) else {
                            eprintln!("Warning: Failed to deserialize secret on line {i}");
                            continue
                        };
                        secrets.push(secret);
                    }
                }

                let pubkeys = match drk.import_money_secrets(secrets).await {
                    Ok(p) => p,
                    Err(e) => {
                        eprintln!("Failed to import secret keys into wallet: {e:?}");
                        exit(2);
                    }
                };

                for key in pubkeys {
                    eprintln!("{key}");
                }

                return Ok(())
            }

            if tree {
                let tree = drk.get_money_tree().await?;

                eprintln!("{tree:#?}");

                return Ok(())
            }

            if coins {
                let coins = drk.get_coins(true).await?;

                let aliases_map = drk.get_aliases_mapped_by_token().await?;

                if coins.is_empty() {
                    return Ok(())
                }

                let mut table = Table::new();
                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
                table.set_titles(row![
                    "Coin",
                    "Spent",
                    "Token ID",
                    "Aliases",
                    "Value",
                    "Spend Hook",
                    "User Data"
                ]);
                let zero = pallas::Base::zero();
                for coin in coins {
                    let aliases = match aliases_map.get(&coin.0.note.token_id.to_string()) {
                        Some(a) => a,
                        None => "-",
                    };

                    let spend_hook = if coin.0.note.spend_hook != zero {
                        bs58::encode(&serialize(&coin.0.note.spend_hook)).into_string().to_string()
                    } else {
                        String::from("-")
                    };

                    let user_data = if coin.0.note.user_data != zero {
                        bs58::encode(&serialize(&coin.0.note.user_data)).into_string().to_string()
                    } else {
                        String::from("-")
                    };

                    table.add_row(row![
                        bs58::encode(&serialize(&coin.0.coin.inner())).into_string().to_string(),
                        coin.1,
                        coin.0.note.token_id,
                        aliases,
                        format!("{} ({})", coin.0.note.value, encode_base10(coin.0.note.value, 8)),
                        spend_hook,
                        user_data
                    ]);
                }

                eprintln!("{table}");

                return Ok(())
            }

            unreachable!()
        }

        Subcmd::Unspend { coin } => {
            let bytes: [u8; 32] = match bs58::decode(&coin).into_vec()?.try_into() {
                Ok(b) => b,
                Err(e) => {
                    eprintln!("Invalid coin: {e:?}");
                    exit(2);
                }
            };

            let elem: pallas::Base = match pallas::Base::from_repr(bytes).into() {
                Some(v) => v,
                None => {
                    eprintln!("Invalid coin");
                    exit(2);
                }
            };

            let coin = Coin::from(elem);
            let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;
            if let Err(e) = drk.unspend_coin(&coin).await {
                eprintln!("Failed to mark coin as unspent: {e:?}");
                exit(2);
            }

            Ok(())
        }

        Subcmd::Inspect => {
            let mut buf = String::new();
            stdin().read_to_string(&mut buf)?;
            let bytes = bs58::decode(&buf.trim()).into_vec()?;
            let tx: Transaction = deserialize(&bytes)?;
            eprintln!("{tx:#?}");
            Ok(())
        }

        Subcmd::Broadcast => {
            eprintln!("Reading transaction from stdin...");
            let mut buf = String::new();
            stdin().read_to_string(&mut buf)?;
            let bytes = bs58::decode(&buf.trim()).into_vec()?;
            let tx = deserialize(&bytes)?;

            let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;

            let txid = match drk.broadcast_tx(&tx).await {
                Ok(t) => t,
                Err(e) => {
                    eprintln!("Failed to broadcast transaction: {e:?}");
                    exit(2);
                }
            };

            eprintln!("Transaction ID: {txid}");

            Ok(())
        }

        Subcmd::Subscribe => {
            let drk =
                Drk::new(args.wallet_path, args.wallet_pass, args.endpoint.clone(), ex.clone())
                    .await?;

            if let Err(e) = drk.subscribe_blocks(args.endpoint, ex).await {
                eprintln!("Block subscription failed: {e:?}");
                exit(2);
            }

            Ok(())
        }

        Subcmd::Scan { reset, list, checkpoint } => {
            let drk =
                Drk::new(args.wallet_path, args.wallet_pass, args.endpoint.clone(), ex.clone())
                    .await?;

            if reset {
                eprintln!("Reset requested.");
                if let Err(e) = drk.scan_blocks(true).await {
                    eprintln!("Failed during scanning: {e:?}");
                    exit(2);
                }
                eprintln!("Finished scanning blockchain");

                return Ok(())
            }

            if list {
                eprintln!("List requested.");
                // TODO: implement

                return Ok(())
            }

            if let Some(c) = checkpoint {
                eprintln!("Checkpoint requested: {c}");
                // TODO: implement

                return Ok(())
            }

            if let Err(e) = drk.scan_blocks(false).await {
                eprintln!("Failed during scanning: {e:?}");
                exit(2);
            }
            eprintln!("Finished scanning blockchain");

            Ok(())
        }

        Subcmd::Explorer { command } => match command {
            ExplorerSubcmd::FetchTx { tx_hash, full, encode } => {
                let tx_hash = blake3::Hash::from_hex(&tx_hash)?;

                let drk =
                    Drk::new(args.wallet_path, args.wallet_pass, args.endpoint.clone(), ex.clone())
                        .await?;

                let tx = match drk.get_tx(&tx_hash).await {
                    Ok(tx) => tx,
                    Err(e) => {
                        eprintln!("Failed to fetch transaction: {e:?}");
                        exit(2);
                    }
                };

                let Some(tx) = tx else {
                    eprintln!("Transaction was not found");
                    exit(1);
                };

                // Make sure the tx is correct
                assert_eq!(tx.hash()?, tx_hash);

                if encode {
                    eprintln!("{}", bs58::encode(&serialize(&tx)).into_string());
                    exit(1)
                }

                eprintln!("Transaction ID: {tx_hash}");
                if full {
                    eprintln!("{tx:?}");
                }

                Ok(())
            }

            ExplorerSubcmd::SimulateTx => {
                eprintln!("Reading transaction from stdin...");
                let mut buf = String::new();
                stdin().read_to_string(&mut buf)?;
                let bytes = bs58::decode(&buf.trim()).into_vec()?;
                let tx = deserialize(&bytes)?;

                let drk =
                    Drk::new(args.wallet_path, args.wallet_pass, args.endpoint.clone(), ex.clone())
                        .await?;

                let is_valid = match drk.simulate_tx(&tx).await {
                    Ok(b) => b,
                    Err(e) => {
                        eprintln!("Failed to simulate tx: {e:?}");
                        exit(2);
                    }
                };

                eprintln!("Transaction ID: {}", tx.hash()?);
                eprintln!("State: {}", if is_valid { "valid" } else { "invalid" });

                Ok(())
            }

            ExplorerSubcmd::TxsHistory { tx_hash, encode } => {
                let drk =
                    Drk::new(args.wallet_path, args.wallet_pass, args.endpoint.clone(), ex.clone())
                        .await?;

                if let Some(c) = tx_hash {
                    let (tx_hash, status, tx) = drk.get_tx_history_record(&c).await?;

                    if encode {
                        println!("{}", bs58::encode(&serialize(&tx)).into_string());
                        exit(1)
                    }

                    eprintln!("Transaction ID: {tx_hash}");
                    eprintln!("Status: {status}");
                    eprintln!("{tx:?}");

                    return Ok(())
                }

                let map = match drk.get_txs_history().await {
                    Ok(m) => m,
                    Err(e) => {
                        eprintln!("Failed to retrieve transactions history records: {e:?}");
                        exit(2);
                    }
                };

                // Create a prettytable with the new data:
                let mut table = Table::new();
                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
                table.set_titles(row!["Transaction Hash", "Status"]);
                for (txs_hash, status) in map.iter() {
                    table.add_row(row![txs_hash, status]);
                }

                if table.is_empty() {
                    eprintln!("No transactions found");
                } else {
                    eprintln!("{table}");
                }

                Ok(())
            }
        },

        Subcmd::Alias { command } => match command {
            AliasSubcmd::Add { alias, token } => {
                if alias.chars().count() > 5 {
                    eprintln!("Error: Alias exceeds 5 characters");
                    exit(2);
                }

                let token_id = match TokenId::from_str(token.as_str()) {
                    Ok(t) => t,
                    Err(e) => {
                        eprintln!("Invalid Token ID: {e:?}");
                        exit(2);
                    }
                };

                let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;
                if let Err(e) = drk.add_alias(alias, token_id).await {
                    eprintln!("Failed to add alias: {e:?}");
                    exit(2);
                }

                Ok(())
            }

            AliasSubcmd::Show { alias, token } => {
                let token_id = match token {
                    Some(t) => match TokenId::from_str(t.as_str()) {
                        Ok(t) => Some(t),
                        Err(e) => {
                            eprintln!("Invalid Token ID: {e:?}");
                            exit(2);
                        }
                    },
                    None => None,
                };

                let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;
                let map = drk.get_aliases(alias, token_id).await?;

                // Create a prettytable with the new data:
                let mut table = Table::new();
                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
                table.set_titles(row!["Alias", "Token ID"]);
                for (alias, token_id) in map.iter() {
                    table.add_row(row![alias, token_id]);
                }

                if table.is_empty() {
                    eprintln!("No aliases found");
                } else {
                    eprintln!("{table}");
                }

                Ok(())
            }

            AliasSubcmd::Remove { alias } => {
                let drk = Drk::new(args.wallet_path, args.wallet_pass, args.endpoint, ex).await?;
                if let Err(e) = drk.remove_alias(alias).await {
                    eprintln!("Failed to remove alias: {e:?}");
                    exit(2);
                }

                Ok(())
            }
        },
    }
}