use nftnl::{nft_expr, Batch, Chain, FinalizedBatch, ProtoFamily, Rule, Table};
use std::{
    ffi::{self, CString},
    io,
};
use mnl;
use libc;

const TABLE_NAME: &str = "diplonat";
const OUT_CHAIN_NAME: &str = "out";
const IN_CHAIN_NAME: &str = "in";

#[derive(Debug)]
struct Error(String);

impl From<io::Error> for Error {
    fn from(error: io::Error) -> Self {
        Error(error.to_string())
    }
}

impl From<ffi::NulError> for Error {
    fn from(error: ffi::NulError) -> Self {
        Error(error.to_string())
    }
}

fn send_and_process(batch: &FinalizedBatch) -> Result<(), Error> {
    let socket = mnl::Socket::new(mnl::Bus::Netfilter)?;
    socket.send_all(batch)?;

    let portid = socket.portid();
    let mut buffer = vec![0; nftnl::nft_nlmsg_maxsize() as usize];

    while let Some(message) = socket_recv(&socket, &mut buffer[..])? {
        match mnl::cb_run(message, 2, portid)? {
            mnl::CbResult::Stop => {
                break;
            }
            mnl::CbResult::Ok => (),
        }
    }
    Ok(())
}

fn socket_recv<'a>(socket: &mnl::Socket, buf: &'a mut [u8]) -> Result<Option<&'a [u8]>, Error> {
    let ret = socket.recv(buf)?;
    if ret > 0 {
        Ok(Some(&buf[..ret]))
    } else {
        Ok(None)
    }
}

fn add_port_allowed(port: u16) -> Result<(), Error> {
    let mut batch = Batch::new();

    // TODO: at the moment, I haven't found a way to properly separate setup part (create table and
    // chains) and the add rule part because the add rule part needs a reference on the chains.
    // apparently creating a table and chains that already exist seems to do nothing so it works
    // doing the following. To be done properly though
    
    let table = Table::new(&CString::new(TABLE_NAME).unwrap(), ProtoFamily::Inet);
    
    batch.add(&table, nftnl::MsgType::Add);

    let mut out_chain = Chain::new(&CString::new(OUT_CHAIN_NAME).unwrap(), &table);
    let mut in_chain = Chain::new(&CString::new(IN_CHAIN_NAME).unwrap(), &table);

    out_chain.set_hook(nftnl::Hook::Out, 0);
    in_chain.set_hook(nftnl::Hook::In, 0);

    out_chain.set_policy(nftnl::Policy::Accept);
    in_chain.set_policy(nftnl::Policy::Drop);

    batch.add(&out_chain, nftnl::MsgType::Add);
    batch.add(&in_chain, nftnl::MsgType::Add);

    let mut _rule = Rule::new(&in_chain);

    _rule.add_expr(&nft_expr!(meta nfproto));
    _rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));

    _rule.add_expr(&nft_expr!(meta l4proto));
    _rule.add_expr(&nft_expr!(cmp == libc::IPPROTO_TCP as u8));

    _rule.add_expr(&nftnl::expr::Payload::Transport(
        nftnl::expr::TransportHeaderField::Tcp(nftnl::expr::TcpHeaderField::Dport),
    ));
    _rule.add_expr(&nft_expr!(cmp == port.to_be()));

    _rule.add_expr(&nft_expr!(verdict accept));

    batch.add(&_rule, nftnl::MsgType::Add);

    let finalized_batch = batch.finalize();
    send_and_process(&finalized_batch)?;
    
    Ok(())
 
}