unshell-nodes-rs/src/main.rs

use std::{
    env,
    error::Error,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    str::FromStr,
};

use clap::{Parser, Subcommand};
use log::error;
use unshell_rs::Cli;
use unshell_rs_lib::connection::Node;
// use unshell_rs::{UnshellClient, UnshellGui, UnshellServer};
// use unshell_rs

pub static DEFAULT_CONFIG_FILEPATH: &'static str = "server_config.json";

// The default port that this program looks for
pub static DEFAULT_SERVICE_PORT: u16 = 13370;
// The default website port that this program looks for
pub static DEFAULT_WEB_PORT: u16 = 8082;

pub static LOCAL_SOCKET: SocketAddr =
    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 12, 34, 56)), 13370);

#[derive(Debug, Parser)]
#[command(name = "unshell-rs")]
#[command(about = "Slick reverse shell tool in rust", long_about = None)]
struct Args {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// Run as a service, and potentially hosting a website
    #[command(arg_required_else_help = true)]
    Relay {
        /// IPv4 to listen for clients on.
        host: String,

        /// Port listen to for command clients
        #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
        port: u16,

        /// Json file to store config
        #[arg(short, long, default_value_t = DEFAULT_CONFIG_FILEPATH.to_string())]
        config_filepath: String,
        // /// Port to listen for website traffic (0 is disabled)
        // #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
        // web_port: u16,
    },
    /// Connect to remote server
    Connect {
        /// Remote server to connect to
        host: String,

        /// Port listen to for command clients
        #[arg(short, long, default_value_t = DEFAULT_SERVICE_PORT)]
        port: u16,
    },
}

fn main() -> Result<(), Box<dyn Error>> {
    if env::var("RUST_LOG").is_err() {
        unsafe { env::set_var("RUST_LOG", "info") }
    }

    pretty_env_logger::init();
    let args = Args::parse();

    match args.command {
        Commands::Relay {
            host,
            port,
            config_filepath,
        } => {
            let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
            if let Err(e) = Node::run(if let Ok(addr) = addr {
                addr
            } else {
                error!("Could not parse address!");
                return Ok(());
            }) {
                error!("{}", e);
            }
        }
        Commands::Connect { host, port } => {
            let addr = SocketAddr::from_str(format!("{}:{}", host, port).as_str());
            if let Err(e) = Cli::connect(if let Ok(addr) = addr {
                addr
            } else {
                error!("Could not parse address!");
                return Ok(());
            }) {
                error!("{}", e);
            }
        }
    };

    Ok(())
}