// sandhole: Expose HTTP/SSH/TCP services through SSH port forwarding
// Copyright (C) 2024-2026 Eric Rodrigues Pires
//
// 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::{sync::Arc, time::Duration};

use bytes::Bytes;
use clap::Parser;

use rand::Rng;
use russh::{
    Channel, ChannelMsg, Preferred,
    client::{ChannelOpenHandle, Msg, Session},
};
use russh::{
    client::Config,
    keys::{key::PrivateKeyWithHashAlg, load_secret_key},
};
use sandhole::{ApplicationConfig, entrypoint};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
    time::{sleep, timeout},
};

use crate::common::SandholeHandle;

/// This test ensures that a TCP service can handle multiple big uploads at the
/// same time (mostly for profiling purposes).
#[test_log::test(tokio::test(flavor = "multi_thread"))]
async fn tcp_multi_stream_upload() {
    // 1. Initialize Sandhole
    let config = ApplicationConfig::parse_from([
        "sandhole",
        "--domain=foobar.tld",
        "--user-keys-directory",
        &(format!(
            "{}/tests/data/user_keys",
            std::env::var("CARGO_MANIFEST_DIR").unwrap()
        )),
        "--admin-keys-directory",
        &(format!(
            "{}/tests/data/admin_keys",
            std::env::var("CARGO_MANIFEST_DIR").unwrap()
        )),
        "--certificates-directory",
        &(format!(
            "{}/tests/data/certificates",
            std::env::var("CARGO_MANIFEST_DIR").unwrap()
        )),
        "--private-key-file",
        &(format!(
            "{}/tests/data/server_keys/ssh",
            std::env::var("CARGO_MANIFEST_DIR").unwrap()
        )),
        "--acme-cache-directory",
        &(format!(
            "{}/tests/data/acme_cache",
            std::env::var("CARGO_MANIFEST_DIR").unwrap()
        )),
        "--disable-directory-creation",
        "--listen-address=127.0.0.1",
        "--ssh-port=18022",
        "--http-port=18080",
        "--https-port=18443",
        "--acme-use-staging",
        "--bind-hostnames=none",
        "--allow-requested-ports",
        "--idle-connection-timeout=1s",
        "--authentication-request-timeout=5s",
        "--ssh-keepalive-interval=45s",
        "--ssh-keepalive-max=2",
    ]);
    let _sandhole_handle = SandholeHandle(tokio::spawn(async move { entrypoint(config).await }));
    if timeout(Duration::from_secs(5), async {
        while TcpStream::connect("127.0.0.1:18022").await.is_err() {
            sleep(Duration::from_millis(100)).await;
        }
    })
    .await
    .is_err()
    {
        panic!("Timeout waiting for Sandhole to start.")
    };

    // 2. Start SSH client that will be proxied
    let key = load_secret_key(
        std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap())
            .join("tests/data/private_keys/key1"),
        None,
    )
    .expect("Missing file key1");
    let mut data = vec![0u8; 20_000_000];
    rand::rng().fill_bytes(&mut data);
    let ssh_client = SshClient(Bytes::from_static(data.leak()));
    let mut session = russh::client::connect(
        Arc::new(Config {
            preferred: Preferred {
                cipher: std::borrow::Cow::Borrowed(&[
                    russh::cipher::CHACHA20_POLY1305,
                    // russh::cipher::AES_256_GCM,
                ]),
                ..Default::default()
            },
            ..Default::default()
        }),
        "127.0.0.1:18022",
        ssh_client,
    )
    .await
    .expect("Failed to connect to SSH server");
    assert!(
        session
            .authenticate_publickey(
                "user",
                PrivateKeyWithHashAlg::new(
                    Arc::new(key),
                    session.best_supported_rsa_hash().await.unwrap().flatten()
                )
            )
            .await
            .expect("SSH authentication failed")
            .success(),
        "authentication didn't succeed"
    );
    session
        .tcpip_forward("tcp.sandhole", 12345)
        .await
        .expect("tcpip_forward failed");

    // 3. Connect to the TCP port of our proxy with out multiple streams
    timeout(Duration::from_secs(120), async move {
        let mut jh_vec = vec![];
        for file_size in [7_500_000usize, 10_000_000, 15_000_000, 20_000_000] {
            let tcp_stream = TcpStream::connect("127.0.0.1:12345")
                .await
                .expect("TCP connection failed");
            tcp_stream.set_nodelay(true).unwrap();
            let (mut read_half, mut write_half) = tcp_stream.into_split();
            let jh = tokio::spawn(async move {
                let jh = tokio::spawn(async move {
                    write_half
                        .write_all(&file_size.to_le_bytes()[..])
                        .await
                        .unwrap();
                });
                let mut buf = vec![0u8; file_size];
                read_half.read_exact(&mut buf).await.unwrap();
                jh.abort();
            });
            jh_vec.push(jh);
        }
        for jh in jh_vec.into_iter() {
            jh.await.expect("Join handle panicked");
        }
    })
    .await
    .expect("Timeout waiting for test to finish.");
}

struct SshClient(Bytes);

impl russh::client::Handler for SshClient {
    type Error = color_eyre::eyre::Error;

    async fn check_server_key(
        &mut self,
        _key: &russh::keys::PublicKey,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }

    async fn server_channel_open_forwarded_tcpip(
        &mut self,
        mut channel: Channel<Msg>,
        _connected_address: &str,
        _connected_port: u32,
        _originator_address: &str,
        _originator_port: u32,
        reply: ChannelOpenHandle,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        let bytes = self.0.clone();
        tokio::spawn(async move {
            let Some(ChannelMsg::Data { data }) = channel.wait().await else {
                panic!("Received invalid message");
            };
            if data.len() == size_of::<usize>() {
                channel
                    .data(
                        &bytes[..usize::from_le_bytes(
                            data[..size_of::<usize>()].try_into().unwrap(),
                        )],
                    )
                    .await
                    .unwrap();
            }
        });
        reply.accept().await;
        Ok(())
    }
}