use std::fs::{OpenOptions, File, create_dir_all};
use std::io::Write;
use std::path::Path;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;

use chrono::Utc;
use futures::FutureExt;
use log::info;
use russh::server::{Auth, Session};
use russh::*;

#[tokio::main]
async fn main() {
    env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .init();

    let listen_addr = std::env::args().skip(1).next()
        .expect("Expecting <addr:port>");

    let mut config = russh::server::Config::default();
    config.auth_banner = Some("#
# DEMONSTRATION SYSTEM
#
# All input will be published
#
");
    config.connection_timeout = Some(std::time::Duration::from_secs(10000));
    config.auth_rejection_time = std::time::Duration::from_secs(1);
    config.keys.push(russh_keys::key::KeyPair::generate_ed25519().unwrap());
    let config = Arc::new(config);
    let sh = Server;
    russh::server::run(
        config,
        &std::net::SocketAddr::from_str(&listen_addr).unwrap(),
        sh,
    )
    .await
    .unwrap();
}

struct Server;

impl server::Server for Server {
    type Handler = Handler;
    fn new_client(&mut self, addr: Option<std::net::SocketAddr>) -> Handler {
        let addr = if let Some(addr) = addr {
            format!("{}", addr)
        } else {
            format!("unknown")
        };
        info!("Connection from {}", addr);

        let path = &format!(
            "{}-{}.txt",
            Utc::now().format("%Y-%m-%d/%H:%M:%S"),
            addr
        );
        let path = Path::new(path);
        create_dir_all(path.parent().unwrap()).unwrap();
        Handler {
            filename: path.display().to_string(),
            lazy_file: None,
            buffer: vec![],
            user: "root".into(),
        }
    }
}

fn send_str(session: &mut Session, channel: ChannelId, s: String) {
    let data = CryptoVec::from(s);
    session.data(channel, data);
}

struct Handler {
    filename: String,
    lazy_file: Option<File>,
    buffer: Vec<u8>,
    user: String,
}

impl Handler {
    fn file(&mut self) -> &mut File {
        if self.lazy_file.is_none() {
            let file = OpenOptions::new()
                .create(true)
                .append(true)
                .open(&self.filename)
                .unwrap();
            self.lazy_file = Some(file);
        }
        self.lazy_file.as_mut().unwrap()
    }

    fn send_prompt(&self, session: &mut Session, channel: ChannelId) {
        send_str(session, channel, format!("{}@fnordister:~$ ", self.user));
    }

    fn handle_command<F: FnMut(&str)>(&self, command: String, mut respond: F) {
        let program_len = command.find(|c: char| c.is_whitespace())
            .unwrap_or(command.len());
        let program = &command[..program_len];

        match program {
            "whoami" => respond(&self.user),
            "id" => respond(&format!("uid=0({}) gid=0(root)", self.user)),
            "uname" => respond("Linux fnordister 5.10.0-10-amd64 #1 SMP Debian 5.10.84-1 (2021-12-08) x86_64 GNU/Linux"),
            "pwd" => respond("/"),
            "ls" => {
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 .");
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 ..");
                respond("drwxr-xr-x  18 root root    18 Jan  4  0000 ...");
            },
            "bash" => {}
            "sh" => {}
            "cd" => {}
            "" => {}
            _ => respond(&format!("{}: command not found", program)),
        }
    }
}

impl server::Handler for Handler {
    type Error = anyhow::Error;
    type FutureAuth =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Auth)>> + Send>>;
    type FutureUnit =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session)>> + Send>>;
    type FutureBool =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session, bool)>> + Send>>;

    fn finished_auth(self, auth: Auth) -> Self::FutureAuth {
        async { Ok((self, auth)) }.boxed()
    }

    fn finished_bool(self, b: bool, s: Session) -> Self::FutureBool {
        async move { Ok((self, s, b)) }.boxed()
    }

    fn finished(self, s: Session) -> Self::FutureUnit {
        async { Ok((self, s)) }.boxed()
    }

    fn shell_request(self, channel: ChannelId, mut session: Session) -> Self::FutureUnit {
        info!("shell_request");
        self.send_prompt(&mut session, channel);

        self.finished(session)
    }

    fn exec_request(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        info!("exec_request");
        writeln!(self.file(), "Execute: {}\n", String::from_utf8_lossy(data))
            .unwrap();

        let line = String::from_utf8_lossy(data).into();
        self.handle_command(line, |response| {
            let mut data = Vec::from(response);
            data.extend_from_slice(b"\r\n");
            session.data(channel, data.into());
        });
        session.close(channel);

        self.finished(session)
    }

    fn subsystem_request(
        mut self,
        _channel: ChannelId,
        name: &str,
        session: Session
    ) -> Self::FutureUnit {
        info!("subsystem_request");
        writeln!(self.file(), "Subsystem requested: {}\n", name)
            .unwrap();

        self.finished(session)
    }

    fn auth_password(mut self, user: &str, password: &str) -> Self::FutureAuth {
        writeln!(self.file(), "Authenticated as {} with {}\n", user, password)
            .unwrap();
        self.user = user.into();
        self.finished_auth(server::Auth::Accept)
    }

    fn data(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        self.file().write(data)
            .unwrap();
        // echo input back
        session.data(channel, data.to_vec().into());

        if self.buffer.len() < 1024 {
            self.buffer.extend_from_slice(data);
        }
        if let Some(newline) = self.buffer.iter().position(|b| *b == b'\r') {
            let rest = self.buffer.split_off(newline + 1);
            let line: String = String::from_utf8_lossy(&self.buffer).into();
            self.buffer = rest;

            session.data(channel, b"\n".to_vec().into());
            self.handle_command(line, |response| {
                let mut data = Vec::from(response);
                data.extend_from_slice(b"\r\n");
                session.data(channel, data.into());
            });
            self.send_prompt(&mut session, channel);
        }

        self.finished(session)
    }
}