//! Implementing a DNS resolver for proxied Tor traffic.

use std::collections::HashMap;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::sync::{Arc, Mutex};

use crate::errors::TunnelResult;
use bimap::BiMap;
use bytes::Bytes;
use dns_message_parser::question::QType;
use dns_message_parser::rr::RR;
use dns_message_parser::RCode;
use futures::stream::StreamExt;
use rand::RngCore;
use smoltcp::socket::udp;
use smoltcp::wire::IpEndpoint;
use smoltcp::{
    socket::udp::PacketBuffer as UdpSocketBuffer, socket::udp::PacketMetadata as UdpPacketMetadata,
    wire::IpAddress, wire::Ipv4Address, wire::Ipv6Address,
};
use tracing::{info, trace, warn};

use crate::socket::UdpSocket;
use crate::{OnionTunnel, TunnelScaffolding};

pub(crate) type LockedDnsCookies = Arc<Mutex<DnsCookies>>;

/// A mapping between dispensed fake IP addresses and hostnames.
/// This mapping is specific to a given isolation key.
#[derive(Default)]
struct IsolatedCookies {
    map_v4: BiMap<Ipv4Address, String>,
    map_v6: BiMap<Ipv6Address, String>,
}

impl IsolatedCookies {
    fn new() -> Self {
        Default::default()
    }

    fn make_for_v4(&mut self, mut hostname: String) -> Ipv4Address {
        // NOTE(eta): is this necessary still?
        if hostname.ends_with('.') {
            hostname.pop();
        }

        // Check if we've already generated one.
        if let Some(a) = self.map_v4.get_by_right(&hostname) {
            return *a;
        }

        // Make a random IPv4 address in the range 10.0.0.0/8.
        let mut ret = Ipv4Address([10, 0, 0, 0]);
        rand::thread_rng().fill_bytes(&mut ret.0[1..]);

        self.map_v4.insert(ret, hostname);
        ret
    }

    fn make_for_v6(&mut self, mut hostname: String) -> Ipv6Address {
        // NOTE(eta): is this necessary still?
        if hostname.ends_with('.') {
            hostname.pop();
        }

        // Check if we've already generated one.
        if let Some(a) = self.map_v6.get_by_right(&hostname) {
            return *a;
        }

        // Make a random IPv6 address in the range fec0::/10.
        // This range is the deprecated "site-local address" range, but it should still
        // be fine to use; it's just a convenient reserved block that's definitely not globally
        // routable.

        let mut ret = Ipv6Address([0xfe, 0xc0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
        rand::thread_rng().fill_bytes(&mut ret.0[2..]);

        self.map_v6.insert(ret, hostname);
        ret
    }

    fn lookup_address(&self, addr: &IpAddress) -> Option<String> {
        match addr {
            IpAddress::Ipv4(a) => self.map_v4.get_by_left(a).cloned(),
            IpAddress::Ipv6(a) => self.map_v6.get_by_left(a).cloned(),
        }
    }
}

pub(crate) struct DnsCookies {
    /// Per-isolation-key isolated store of cookies.
    cookies: HashMap<Option<u64>, IsolatedCookies>,
}

impl DnsCookies {
    fn new() -> Self {
        Self {
            cookies: HashMap::new(),
        }
    }

    fn make_for_v4(&mut self, isolation_key: Option<u64>, hostname: String) -> Ipv4Address {
        self.cookies
            .entry(isolation_key)
            .or_insert_with(IsolatedCookies::new)
            .make_for_v4(hostname)
    }

    fn make_for_v6(&mut self, isolation_key: Option<u64>, hostname: String) -> Ipv6Address {
        self.cookies
            .entry(isolation_key)
            .or_insert_with(IsolatedCookies::new)
            .make_for_v6(hostname)
    }

    pub fn get(&self, isolation_key: Option<u64>, addr: &IpAddress) -> Option<String> {
        self.cookies
            .get(&isolation_key)
            .and_then(|x| x.lookup_address(addr))
    }
}

pub struct DnsManager<S> {
    listener_udp_v4: UdpSocket,
    listener_udp_v6: UdpSocket,
    cookies: LockedDnsCookies,
    scaffolding: Arc<S>,
}

impl<S> DnsManager<S>
where
    S: TunnelScaffolding,
{
    /// The IPv4 address the DNS server listens on.
    ///
    /// Uses the 169.254.0.0/16 link-local block.
    const V4_ADDRESS: Ipv4Addr = Ipv4Addr::new(169, 254, 42, 53);
    /// The IPv6 address the DNS server listens on.
    ///
    /// Uses the fe80::/10 link-local block.
    const V6_ADDRESS: Ipv6Addr = Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 53);

    pub fn new<D>(tunnel: &OnionTunnel<S, D>) -> TunnelResult<Self> {
        Ok(Self {
            listener_udp_v4: Self::new_udp_listener(tunnel, Self::V4_ADDRESS.into())?,
            listener_udp_v6: Self::new_udp_listener(tunnel, Self::V6_ADDRESS.into())?,
            cookies: Arc::new(Mutex::new(DnsCookies::new())),
            scaffolding: tunnel.scaffolding.clone(),
        })
    }

    fn new_udp_listener<D>(
        tunnel: &OnionTunnel<S, D>,
        bind_addr: IpAddress,
    ) -> TunnelResult<UdpSocket> {
        let rx_buffer = UdpSocketBuffer::new(vec![UdpPacketMetadata::EMPTY], vec![0; 4096]);
        let tx_buffer = UdpSocketBuffer::new(vec![UdpPacketMetadata::EMPTY], vec![0; 4096]);

        let socket = udp::Socket::new(rx_buffer, tx_buffer);

        let mut listener = UdpSocket::new(tunnel, socket);
        #[allow(clippy::expect_used)]
        listener
            .bind(bind_addr, 53)
            .expect("bind error when creating DNS listener?");

        Ok(listener)
    }

    pub(crate) fn cookies(&self) -> Arc<Mutex<DnsCookies>> {
        self.cookies.clone()
    }

    fn make_answer(
        query: &dns_message_parser::Dns,
        rcode: RCode,
        answers: Vec<RR>,
    ) -> dns_message_parser::Dns {
        let flags = dns_message_parser::Flags {
            qr: true,
            opcode: dns_message_parser::Opcode::Query,
            aa: false,
            tc: false,
            rd: true,
            ra: true,
            ad: false,
            cd: false,
            rcode,
        };
        // NOTE: We can't really add authorities or additionals since we don't know
        //       what they are!
        dns_message_parser::Dns {
            id: query.id,
            flags,
            questions: query.questions.clone(),
            answers,
            authorities: vec![],
            additionals: vec![],
        }
    }

    fn build_answer(
        &self,
        query: &dns_message_parser::Dns,
        from: IpEndpoint,
    ) -> Option<dns_message_parser::Dns> {
        let dest = if matches!(from.addr, IpAddress::Ipv4(..)) {
            IpEndpoint::new(Self::V4_ADDRESS.into(), 53)
        } else {
            IpEndpoint::new(Self::V6_ADDRESS.into(), 53)
        };

        // We try and get the isolation key, in order to isolate DNS per-app.
        // If we can't, we store it in the fallback key, which the proxying code will check if it
        // can't find the isolation key.
        // This does kind of ruin the feature slightly, but it doesn't work properly on some
        // Android platforms without this. Sigh.
        let isolation_key = match self.scaffolding.isolate(from, dest, 17 /* UDP */) {
            Ok(k) => Some(k),
            Err(e) => {
                trace!(
                    "DNS query from {} failed due to isolate() call failure: {}",
                    from,
                    e
                );
                None
                /*
                NOTE(eta): It would be nice if we could do this...
                return Some(Self::make_answer(query, RCode::ServFail, vec![]));
                 */
            }
        };

        let mut answers = vec![];

        for question in query.questions.iter() {
            let name = question.domain_name.to_string();
            let rr = match &question.q_type {
                &QType::A => {
                    #[allow(clippy::expect_used)]
                    let cookie = self
                        .cookies
                        .lock()
                        .expect("cookie poisoned")
                        .make_for_v4(isolation_key, name);
                    RR::A(dns_message_parser::rr::A {
                        domain_name: question.domain_name.clone(),
                        ttl: 3600,
                        ipv4_addr: Ipv4Addr::from(cookie),
                    })
                }
                &QType::AAAA => {
                    #[allow(clippy::expect_used)]
                    let cookie = self
                        .cookies
                        .lock()
                        .expect("cookie poisoned")
                        .make_for_v6(isolation_key, name);
                    RR::AAAA(dns_message_parser::rr::AAAA {
                        domain_name: question.domain_name.clone(),
                        ttl: 3600,
                        ipv6_addr: Ipv6Addr::from(cookie),
                    })
                }
                // TODO(eta): implement PTR
                x => {
                    warn!(
                        "Received DNS query of type {}, which we don't support yet!",
                        x
                    );
                    // return "not implemented" -- at least it's better than nothing!
                    return Some(Self::make_answer(query, RCode::NotImp, vec![]));
                }
            };
            answers.push(rr);
        }

        Some(Self::make_answer(query, RCode::NoError, answers))
    }

    fn decode(&self, payload: Bytes, from: IpEndpoint) -> Option<dns_message_parser::Dns> {
        let ret = dns_message_parser::Dns::decode(payload);
        match ret {
            Ok(p) => {
                trace!("DNS question from {}: {}", from, p);
                let ret = self.build_answer(&p, from);
                if let Some(ref r) = ret {
                    trace!("DNS answer to {}: {}", from, r);
                } else {
                    trace!("no DNS answer for {}", from);
                }
                ret
            }
            Err(e) => {
                warn!("Unable to parse DNS request: {}", e);
                None
            }
        }
    }

    pub async fn start(&mut self) {
        info!("Starting DNS manager");

        loop {
            tokio::select! {
                r = self.listener_udp_v4.next() => match r {
                    Some((payload, endpoint)) => {
                        if let Some(reply) = self.decode(payload, endpoint) {
                            match reply.encode() {
                                Ok(v) => {
                                    if let Err(e) = self.listener_udp_v4.send(&v, endpoint) {
                                        warn!("Failed to send DNS reply: {e:?}");
                                    }
                                },
                                Err(e) => warn!("Failed to encode DNS reply: {e}")
                            }
                        }
                    }
                    None => break,
                },
                r = self.listener_udp_v6.next() => match r {
                    Some((payload, endpoint)) => {
                        if let Some(reply) = self.decode(payload, endpoint) {
                            match reply.encode() {
                                Ok(v) => {
                                    if let Err(e) = self.listener_udp_v6.send(&v, endpoint) {
                                        warn!("Failed to send DNS reply: {e:?}");
                                    }
                                },
                                Err(e) => warn!("Failed to encode DNS reply: {e}")
                            }
                        }
                    }
                    None => break,
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::dns::DnsCookies;
    use smoltcp::wire::{IpAddress, Ipv4Address};

    #[test]
    fn dns_cookie_basics() {
        let mut cookies = DnsCookies::new();

        let hostname_4: String = "test.hostname".into();
        let hostname_6: String = "lol.arpa".into();

        let v4 = IpAddress::Ipv4(cookies.make_for_v4(Some(0), hostname_4.clone()));
        let v6 = IpAddress::Ipv6(cookies.make_for_v6(Some(0), hostname_6.clone()));

        assert_eq!(cookies.get(Some(0), &v4), Some(hostname_4));
        assert_eq!(cookies.get(Some(0), &v6), Some(hostname_6));
        assert_eq!(
            cookies.get(Some(0), &IpAddress::Ipv4(Ipv4Address([1, 2, 3, 4]))),
            None
        );
    }

    #[test]
    fn dns_cookie_addresses_stay_the_same() {
        let mut cookies = DnsCookies::new();

        let hostname_4: String = "test.hostname".into();
        let hostname_6: String = "lol.arpa".into();

        let v4 = IpAddress::Ipv4(cookies.make_for_v4(Some(0), hostname_4.clone()));
        let v6 = IpAddress::Ipv6(cookies.make_for_v6(Some(0), hostname_6.clone()));
        let same_v4 = IpAddress::Ipv4(cookies.make_for_v4(Some(0), hostname_4));
        let same_v6 = IpAddress::Ipv6(cookies.make_for_v6(Some(0), hostname_6));
        assert_eq!(v4, same_v4);
        assert_eq!(v6, same_v6);
    }

    #[test]
    fn dns_cookies_are_isolated() {
        let mut cookies = DnsCookies::new();

        let hostname_4: String = "test.hostname".into();
        let hostname_6: String = "lol.arpa".into();

        let v4_0 = IpAddress::Ipv4(cookies.make_for_v4(Some(0), hostname_4.clone()));
        let v6_0 = IpAddress::Ipv6(cookies.make_for_v6(Some(0), hostname_6.clone()));

        assert_eq!(cookies.get(Some(0), &v4_0), Some(hostname_4));
        assert_eq!(cookies.get(Some(1), &v4_0), None);
        assert_eq!(cookies.get(Some(0), &v6_0), Some(hostname_6));
        assert_eq!(cookies.get(Some(1), &v6_0), None);
    }

    #[test]
    fn dns_cookies_are_distinct_v4() {
        let mut cookies = DnsCookies::new();
        let addresses = (0..64)
            .map(|x| {
                IpAddress::Ipv4(cookies.make_for_v4(Some(0), format!("random-hostname{x}.arpa")))
            })
            .collect::<Vec<_>>();

        // Yes, this is silly and inefficient. No, I don't care much.
        for (ia, a) in addresses.iter().enumerate() {
            for (ib, b) in addresses.iter().enumerate() {
                if ia != ib {
                    assert_ne!(a, b);
                }
            }
        }
    }

    #[test]
    fn dns_cookies_are_distinct_v6() {
        let mut cookies = DnsCookies::new();
        let addresses = (0..64)
            .map(|x| {
                IpAddress::Ipv6(cookies.make_for_v6(Some(0), format!("random-hostname{x}.arpa")))
            })
            .collect::<Vec<_>>();

        // Yes, this is silly and inefficient. No, I don't care much.
        for (ia, a) in addresses.iter().enumerate() {
            for (ib, b) in addresses.iter().enumerate() {
                if ia != ib {
                    assert_ne!(a, b);
                }
            }
        }
    }
}
