lssa/wallet/src/cli/token_program.rs

use anyhow::Result;
use clap::Subcommand;
use common::transaction::NSSATransaction;
use nssa::Address;

use crate::{SubcommandReturnValue, WalletCore, cli::WalletSubcommand};

///Represents generic CLI subcommand for a wallet working with token_program
#[derive(Subcommand, Debug, Clone)]
pub enum TokenProgramSubcommand {
    ///Public execution
    #[command(subcommand)]
    Public(TokenProgramSubcommandPublic),
    ///Private execution
    #[command(subcommand)]
    Private(TokenProgramSubcommandPrivate),
}

///Represents generic public CLI subcommand for a wallet working with token_program
#[derive(Subcommand, Debug, Clone)]
pub enum TokenProgramSubcommandPublic {
    //Create a new token using the token program
    CreateNewToken {
        #[arg(short, long)]
        definition_addr: String,
        #[arg(short, long)]
        supply_addr: String,
        #[arg(short, long)]
        name: String,
        #[arg(short, long)]
        total_supply: u128,
    },
    //Transfer tokens using the token program
    TransferToken {
        #[arg(short, long)]
        sender_addr: String,
        #[arg(short, long)]
        recipient_addr: String,
        #[arg(short, long)]
        balance_to_move: u128,
    },
}

///Represents generic public CLI subcommand for a wallet working with token_program
#[derive(Subcommand, Debug, Clone)]
pub enum TokenProgramSubcommandPrivate {
    //Create a new token using the token program
    CreateNewTokenPrivateOwned {
        #[arg(short, long)]
        definition_addr: String,
        #[arg(short, long)]
        supply_addr: String,
        #[arg(short, long)]
        name: String,
        #[arg(short, long)]
        total_supply: u128,
    },
    //Transfer tokens using the token program
    TransferTokenPrivateOwned {
        #[arg(short, long)]
        sender_addr: String,
        #[arg(short, long)]
        recipient_addr: String,
        #[arg(short, long)]
        balance_to_move: u128,
    },
    //Transfer tokens using the token program
    TransferTokenPrivateForeign {
        #[arg(short, long)]
        sender_addr: String,
        ///recipient_npk - valid 32 byte hex string
        #[arg(long)]
        recipient_npk: String,
        ///recipient_ipk - valid 33 byte hex string
        #[arg(long)]
        recipient_ipk: String,
        #[arg(short, long)]
        balance_to_move: u128,
    },
}

impl WalletSubcommand for TokenProgramSubcommandPublic {
    async fn handle_subcommand(
        self,
        wallet_core: &mut WalletCore,
    ) -> Result<SubcommandReturnValue> {
        match self {
            TokenProgramSubcommandPublic::CreateNewToken {
                definition_addr,
                supply_addr,
                name,
                total_supply,
            } => {
                let name = name.as_bytes();
                if name.len() > 6 {
                    // TODO: return error
                    panic!();
                }
                let mut name_bytes = [0; 6];
                name_bytes[..name.len()].copy_from_slice(name);
                wallet_core
                    .send_new_token_definition(
                        definition_addr.parse().unwrap(),
                        supply_addr.parse().unwrap(),
                        name_bytes,
                        total_supply,
                    )
                    .await?;
                Ok(SubcommandReturnValue::Empty)
            }
            TokenProgramSubcommandPublic::TransferToken {
                sender_addr,
                recipient_addr,
                balance_to_move,
            } => {
                wallet_core
                    .send_transfer_token_transaction(
                        sender_addr.parse().unwrap(),
                        recipient_addr.parse().unwrap(),
                        balance_to_move,
                    )
                    .await?;
                Ok(SubcommandReturnValue::Empty)
            }
        }
    }
}

impl WalletSubcommand for TokenProgramSubcommandPrivate {
    async fn handle_subcommand(
        self,
        wallet_core: &mut WalletCore,
    ) -> Result<SubcommandReturnValue> {
        match self {
            TokenProgramSubcommandPrivate::CreateNewTokenPrivateOwned {
                definition_addr,
                supply_addr,
                name,
                total_supply,
            } => {
                let name = name.as_bytes();
                if name.len() > 6 {
                    // TODO: return error
                    panic!("Name length mismatch");
                }
                let mut name_bytes = [0; 6];
                name_bytes[..name.len()].copy_from_slice(name);

                let definition_addr: Address = definition_addr.parse().unwrap();
                let supply_addr: Address = supply_addr.parse().unwrap();

                let (res, [secret_supply]) = wallet_core
                    .send_new_token_definition_private_owned(
                        definition_addr,
                        supply_addr,
                        name_bytes,
                        total_supply,
                    )
                    .await?;

                println!("Results of tx send is {res:#?}");

                let tx_hash = res.tx_hash;
                let transfer_tx = wallet_core
                    .poll_native_token_transfer(tx_hash.clone())
                    .await?;

                if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx {
                    let acc_decode_data = vec![(secret_supply, supply_addr)];

                    wallet_core.decode_insert_privacy_preserving_transaction_results(
                        tx,
                        &acc_decode_data,
                    )?;
                }

                let path = wallet_core.store_persistent_accounts()?;

                println!("Stored persistent accounts at {path:#?}");

                Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
            }
            TokenProgramSubcommandPrivate::TransferTokenPrivateOwned {
                sender_addr,
                recipient_addr,
                balance_to_move,
            } => {
                let sender_addr: Address = sender_addr.parse().unwrap();
                let recipient_addr: Address = recipient_addr.parse().unwrap();

                let recipient_initialization = wallet_core
                    .check_private_account_initialized(&recipient_addr)
                    .await?;

                let (res, [secret_sender, secret_recipient]) =
                    if let Some(recipient_proof) = recipient_initialization {
                        wallet_core
                        .send_transfer_token_transaction_private_owned_account_already_initialized(
                            sender_addr,
                            recipient_addr,
                            balance_to_move,
                            recipient_proof,
                        )
                        .await?
                    } else {
                        wallet_core
                            .send_transfer_token_transaction_private_owned_account_not_initialized(
                                sender_addr,
                                recipient_addr,
                                balance_to_move,
                            )
                            .await?
                    };

                println!("Results of tx send is {res:#?}");

                let tx_hash = res.tx_hash;
                let transfer_tx = wallet_core
                    .poll_native_token_transfer(tx_hash.clone())
                    .await?;

                if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx {
                    let acc_decode_data = vec![
                        (secret_sender, sender_addr),
                        (secret_recipient, recipient_addr),
                    ];

                    wallet_core.decode_insert_privacy_preserving_transaction_results(
                        tx,
                        &acc_decode_data,
                    )?;
                }

                let path = wallet_core.store_persistent_accounts()?;

                println!("Stored persistent accounts at {path:#?}");

                Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
            }
            TokenProgramSubcommandPrivate::TransferTokenPrivateForeign {
                sender_addr,
                recipient_npk,
                recipient_ipk,
                balance_to_move,
            } => {
                let sender_addr: Address = sender_addr.parse().unwrap();
                let recipient_npk_res = hex::decode(recipient_npk)?;
                let mut recipient_npk = [0; 32];
                recipient_npk.copy_from_slice(&recipient_npk_res);
                let recipient_npk = nssa_core::NullifierPublicKey(recipient_npk);

                let recipient_ipk_res = hex::decode(recipient_ipk)?;
                let mut recipient_ipk = [0u8; 33];
                recipient_ipk.copy_from_slice(&recipient_ipk_res);
                let recipient_ipk = nssa_core::encryption::shared_key_derivation::Secp256k1Point(
                    recipient_ipk.to_vec(),
                );

                let (res, [secret_sender, _]) = wallet_core
                    .send_transfer_token_transaction_private_foreign_account(
                        sender_addr,
                        recipient_npk,
                        recipient_ipk,
                        balance_to_move,
                    )
                    .await?;

                println!("Results of tx send is {res:#?}");

                let tx_hash = res.tx_hash;
                let transfer_tx = wallet_core
                    .poll_native_token_transfer(tx_hash.clone())
                    .await?;

                if let NSSATransaction::PrivacyPreserving(tx) = transfer_tx {
                    let acc_decode_data = vec![(secret_sender, sender_addr)];

                    wallet_core.decode_insert_privacy_preserving_transaction_results(
                        tx,
                        &acc_decode_data,
                    )?;
                }

                let path = wallet_core.store_persistent_accounts()?;

                println!("Stored persistent accounts at {path:#?}");

                Ok(SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash })
            }
        }
    }
}

impl WalletSubcommand for TokenProgramSubcommand {
    async fn handle_subcommand(
        self,
        wallet_core: &mut WalletCore,
    ) -> Result<SubcommandReturnValue> {
        match self {
            TokenProgramSubcommand::Private(private_subcommand) => {
                private_subcommand.handle_subcommand(wallet_core).await
            }
            TokenProgramSubcommand::Public(public_subcommand) => {
                public_subcommand.handle_subcommand(wallet_core).await
            }
        }
    }
}