use anyhow::Context;
use clap::{Arg, Command};
use libc::{
    __c_anonymous_ifr_ifru, c_char, ifreq, in_addr, rtentry, sockaddr, sockaddr_in, AF_INET,
    IFF_RUNNING, IFF_UP, INADDR_ANY, RTF_GATEWAY, RTF_UP, SIGCHLD, SIOCADDRT, SIOCGIFFLAGS,
    SIOCSIFADDR, SIOCSIFFLAGS,
};
use log::{info, trace, warn};
use nix::cmsg_space;
use nix::sched::{clone, CloneFlags};
use nix::sys::socket::{
    recvmsg, sendmsg, socketpair, AddressFamily, ControlMessage, ControlMessageOwned, MsgFlags,
    SockFlag, SockType,
};
use nix::unistd::{execv, getgid, getuid, setuid, Gid, Uid};
use onion_tunnel::scaffolding::{ConnectionDetails, FailedConnectionDetails};
use onion_tunnel::{CountryCode, IpEndpoint, OnionTunnel, TunnelScaffolding};
use simple_proc_net::ProcNetEntry;
use smoltcp::phy::{Medium, TunTapInterface};
use std::ffi::{c_void, CString};
use std::fs::File;
use std::mem::size_of;
use std::net::{Ipv4Addr, SocketAddr};
use std::os::fd::RawFd;
use std::os::fd::{AsRawFd, OwnedFd};
use std::str::FromStr;
use std::{fs, io, mem};
use tokio::net::TcpSocket;
use tracing_subscriber::FmtSubscriber;

fn subshell_main(fd: OwnedFd, uid: Uid, gid: Gid) -> isize {
    eprintln!("subshell: setting up environment...");
    fs::write("/proc/self/uid_map", format!("0 {} 1", uid.as_raw())).unwrap();
    // In the case of gid_map, use of the setgroups(2) system
    // call must first be denied by writing "deny" to the
    // /proc/pid/setgroups file (see below) before writing to
    // gid_map.
    fs::write("/proc/self/setgroups", "deny").unwrap();
    fs::write("/proc/self/gid_map", format!("0 {} 1", gid.as_raw())).unwrap();
    eprintln!("subshell: creating tun/tap device...");
    let tti = TunTapInterface::new("onion0", Medium::Ip).unwrap();
    let tti_fd = tti.as_raw_fd();
    mem::forget(tti);
    let mut ifr_name: [c_char; 16] = [0; 16];
    for (i, &b) in "onion0".as_bytes().iter().enumerate() {
        ifr_name[i] = b as c_char;
    }
    let mut ifr = ifreq {
        ifr_name,
        ifr_ifru: __c_anonymous_ifr_ifru { ifru_flags: 0 },
    };
    let socket = TcpSocket::new_v4().unwrap();
    let ret = unsafe { libc::ioctl(socket.as_raw_fd(), SIOCGIFFLAGS, &ifr as *const ifreq) };
    if ret < 0 {
        panic!(
            "ioctl(SIOCGIFFLAGS) failed: {}",
            io::Error::from_raw_os_error(ret)
        );
    }
    unsafe {
        eprintln!("flags: {}", ifr.ifr_ifru.ifru_flags);
        ifr.ifr_ifru.ifru_flags |= (IFF_UP | IFF_RUNNING) as i16;
    }
    let ret = unsafe { libc::ioctl(socket.as_raw_fd(), SIOCSIFFLAGS, &ifr as *const ifreq) };
    if ret < 0 {
        panic!(
            "ioctl(SIOCSIFFLAGS) failed: {}",
            io::Error::from_raw_os_error(ret)
        );
    }
    let address: Ipv4Addr = "169.254.42.1".parse().unwrap();
    let sockaddr = sockaddr_in {
        sin_family: AF_INET as _,
        sin_port: 0,
        sin_addr: in_addr {
            s_addr: u32::from_ne_bytes(address.octets()),
        },
        sin_zero: [0; 8],
    };
    let ret = unsafe {
        ifr.ifr_ifru.ifru_addr = mem::transmute(sockaddr);
        libc::ioctl(socket.as_raw_fd(), SIOCSIFADDR, &ifr as *const ifreq)
    };
    if ret < 0 {
        panic!("ioctl(SIOCSIFADDR) failed: {}", io::Error::last_os_error());
    }
    let device = CString::new("onion0").unwrap();
    let rte = unsafe {
        let mut rte = mem::zeroed::<rtentry>();
        let rt_dst = mem::transmute::<_, &mut sockaddr_in>(&mut rte.rt_dst);
        rt_dst.sin_family = AF_INET as _;
        rt_dst.sin_addr.s_addr = INADDR_ANY;
        let rt_genmask = mem::transmute::<_, &mut sockaddr_in>(&mut rte.rt_genmask);
        rt_genmask.sin_family = AF_INET as _;
        rt_genmask.sin_addr.s_addr = INADDR_ANY;
        rte.rt_flags = RTF_UP;
        rte.rt_metric = 0;
        rte.rt_dev = device.as_ptr() as *mut _;
        rte
    };
    let ret = unsafe { libc::ioctl(socket.as_raw_fd(), SIOCADDRT, &rte as *const rtentry) };
    if ret < 0 {
        panic!("ioctl(SIOCADDRT) failed: {}", io::Error::last_os_error());
    }
    eprintln!("subshell: transmitting tun/tap fd to parent...");
    let tti_fds = [tti_fd];
    let cmsg = ControlMessage::ScmRights(&tti_fds);
    sendmsg::<()>(fd.as_raw_fd(), &[], &[cmsg], MsgFlags::empty(), None).unwrap();
    let path = CString::new("/bin/bash").unwrap();
    let _ = execv::<CString>(&path, &[]);
    0
}

struct LinuxScaffolding {
    cc: Option<CountryCode>,
    can_mark: bool,
    log_connections: bool,
}

impl LinuxScaffolding {
    /// The fwmark to set on arti connections (so policy routing can avoid them being routed
    /// recursively back into the tunnel).
    ///
    /// This is just a randomly generated set of 2 bytes that hopefully won't conflict with
    /// anything else.
    pub const FWMARK: libc::c_int = 0xc185;

    /// Mark the provided socket file descriptor with the `FWMARK`.
    #[cfg(target_os = "linux")]
    fn mark_fd(fd: RawFd) -> io::Result<()> {
        let ret = unsafe {
            libc::setsockopt(
                fd,
                libc::SOL_SOCKET,
                libc::SO_MARK,
                &Self::FWMARK as *const libc::c_int as *const c_void,
                size_of::<libc::c_int>() as _,
            )
        };
        if ret != 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }
}

impl TunnelScaffolding for LinuxScaffolding {
    fn protect(&self, fd: RawFd, _: &SocketAddr) -> io::Result<()> {
        #[cfg(target_os = "linux")]
        if self.can_mark {
            Self::mark_fd(fd)?;
        }
        Ok(())
    }

    fn locate(&self, _: IpEndpoint, _: IpEndpoint, _: u64) -> Option<CountryCode> {
        self.cc
    }

    fn on_bootstrapped(&self) {
        info!("Connection to Tor complete!");
    }

    fn on_established(&self, details: ConnectionDetails<'_>) {
        if self.log_connections {
            let exit_identity = details
                .circuit_relays()
                .last()
                .map(|x| {
                    let cc = match x.country_code {
                        Some(v) => v.to_string(),
                        None => "??".into(),
                    };
                    if let Some(i) = x.ed_identity {
                        format!("{} ({})", i, cc)
                    } else if let Some(i) = x.rsa_identity {
                        format!("{} ({})", i, cc)
                    } else {
                        "???".into()
                    }
                })
                .unwrap_or_else(|| "???".into());

            info!(
                "New connection to '{}' (uid {}) via exit {}",
                details.tor_dst, details.isolation_key, exit_identity
            );
        }
    }

    fn on_arti_failure(&self, details: FailedConnectionDetails) {
        if self.log_connections {
            warn!(
                "Failed to connect to '{}': {}",
                details.tor_dst, details.error
            );
        }
    }

    fn isolate(&self, src: IpEndpoint, dst: IpEndpoint, ip_proto: u8) -> io::Result<u64> {
        let iter = match ip_proto {
            6 => ProcNetEntry::tcp4()?.chain(ProcNetEntry::tcp6()?),
            17 => ProcNetEntry::udp4()?.chain(ProcNetEntry::udp6()?),
            x => panic!("unknown ip_proto {x}"),
        };
        for entry in iter {
            match entry? {
                Ok(v) => {
                    if IpEndpoint::from(v.local_addr) == src
                        && IpEndpoint::from(v.remote_addr) == dst
                    {
                        trace!("isolated {src} -> {dst} proto {ip_proto} as {}", v.uid);
                        return Ok(v.uid as u64);
                    }
                }
                Err(e) => {
                    warn!("Failed to parse /proc line: {}", e.offending_line);
                    warn!("Error encountered: {}", e.error);
                }
            }
        }
        warn!("no isolation found for {src} -> {dst} proto {ip_proto}");
        Ok(0)
    }
}

const ENV_FILTER_VERBOSE: &str =
    "info,smoltcp=debug,onion_tunnel=trace,arti_client=debug,tor_chanmgr=debug,tor_proto=debug";
const ENV_FILTER_TAME: &str = "info";

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let matches = Command::new("onionmasq")
        .version(env!("CARGO_PKG_VERSION"))
        .author("eta <eta@torproject.org>")
        .about("A magical TUN device that feeds traffic via Tor.")
        .arg(
            Arg::new("subshell")
                .short('s')
                .long("subshell")
                .action(clap::ArgAction::SetTrue)
                .help(
                    "Spawn a subshell interactively in a new user namespace with just the onionmasq tun device."
                ),
        )
        .arg(
            Arg::new("tun-device")
                .short('d')
                .long("device")
                .value_name("DEVICE")
                .default_value("onion0")
                .help(
                    "Name of a TUN device to use. Will attempt to create one if it doesn't exist.",
                ),
        )
        .arg(
            Arg::new("country-code")
                .short('c')
                .long("country-code")
                .value_name("DE|NL|etc")
                .help("Make traffic come from exit nodes in this country (ISO 3166-1 alpha-2 country code)"),
        )
        /* TODO(eta): make this a thing
        .arg(Arg::new("whole-system")
            .short('s')
            .long("whole-system")
            .help("Make all system traffic go through the onionmasq TUN device. Off by default.")
            .action(clap::ArgAction::SetTrue)
        )
         */
        .arg(Arg::new("debug")
            .long("debug")
            .help("Enable detailed debug logging. Off by default.")
            .action(clap::ArgAction::SetTrue)
        )
        .arg(Arg::new("verbose")
            .short('v')
            .long("verbose")
            .help("Print information about successful and failed connections. Off by default.")
            .action(clap::ArgAction::SetTrue))
        .get_matches();

    let filter = if matches.get_flag("debug") {
        ENV_FILTER_VERBOSE
    } else {
        ENV_FILTER_TAME
    };

    let log_connections = matches.get_flag("verbose");

    let country_code = matches
        .get_one("country-code")
        .map(|x: &String| CountryCode::from_str(x))
        .transpose()
        .context("invalid country code provided")?;

    let tun_device: &String = matches.get_one("tun-device").unwrap();

    FmtSubscriber::builder().with_env_filter(filter).init();

    let tun_fd = if matches.get_flag("subshell") {
        info!(
            "Starting onionmasq {} with namespaced subshell...",
            env!("CARGO_PKG_VERSION"),
        );
        let uid = getuid();
        let gid = getgid();
        let (parent_fd, child_fd) = socketpair(
            AddressFamily::Unix,
            SockType::Datagram,
            None,
            SockFlag::SOCK_CLOEXEC,
        )
        .unwrap();
        let mut stack = Box::new([0u8; 1024 * 1024 * 2]);
        let mut child_fd = Some(child_fd);
        let pid = unsafe {
            clone(
                Box::new(move || subshell_main(child_fd.take().unwrap(), uid, gid)),
                &mut *stack,
                CloneFlags::CLONE_NEWNET | CloneFlags::CLONE_NEWUSER,
                Some(SIGCHLD),
            )
            .unwrap()
        };
        info!("Spawned child as PID {pid}, waiting for tun/tap fd...");
        let mut space = cmsg_space!([RawFd; 1]);
        let recv = recvmsg::<()>(
            parent_fd.as_raw_fd(),
            &mut [],
            Some(&mut space),
            MsgFlags::empty(),
        )
        .unwrap();
        let cm = recv.cmsgs().next().expect("no control message");
        let tun_fd = if let ControlMessageOwned::ScmRights(fd) = cm {
            fd.into_iter().next().unwrap()
        } else {
            panic!("random control message: {cm:?}");
        };
        info!("Got fd {tun_fd}!");
        Some(tun_fd)
    } else {
        info!(
            "Starting onionmasq {} on device '{}'...",
            env!("CARGO_PKG_VERSION"),
            tun_device
        );
        None
    };

    let can_mark = if matches.get_flag("subshell") {
        false
    } else {
        // Check whether we can call setsockopt(). If we can't, don't bother doing so in future, since
        // we probably don't have the capabilities; just print an explanatory warning instead.
        #[cfg(target_os = "linux")]
        let dummy_socket = TcpSocket::new_v4()?;
        #[cfg(target_os = "linux")]
        match LinuxScaffolding::mark_fd(dummy_socket.as_raw_fd()) {
            Ok(_) => {
                info!(
                    "Outgoing connections (to the Tor network) will use fwmark {:#x}.",
                    LinuxScaffolding::FWMARK
                );
                true
            }
            Err(e) => {
                warn!("Calling setsockopt() failed: {e}");
                warn!("Setting fwmarks on outgoing sockets will be disabled.");
                warn!("Make the onionmasq binary CAP_NET_ADMIN to fix this problem.");
                false
            }
        }

        #[cfg(not(target_os = "linux"))]
        false
    };

    let scaffolding = LinuxScaffolding {
        can_mark,
        cc: country_code,
        log_connections,
    };

    #[cfg(target_os = "linux")]
    let mut onion_tunnel = if let Some(fd) = tun_fd {
        OnionTunnel::create_with_fd(scaffolding, fd, Default::default()).await?
    } else {
        OnionTunnel::new(scaffolding, tun_device, Default::default()).await?
    };

    info!("Connecting to Tor...");

    #[cfg(target_os = "linux")]
    tokio::select! {
        _ = onion_tunnel.run() => (),
        _ = tokio::signal::ctrl_c() => {
            info!("Received ^C; stopping tunnel.");
        }
    };

    Ok(())
}
