use std::net::UdpSocket;
use crate::dns::*;
use std::str;
use crate::utilities::*;


pub struct DnsServer {
}

impl DnsServer {
    pub fn serve(port: u16) {
        let socket = UdpSocket::bind(("0.0.0.0", port)).unwrap();

        // starting main loop
        loop {
            DnsServer::handle_query(&socket);
        }

    }

    fn handle_query(socket: &UdpSocket) {

        let mut response_buffer = BytePacketBuffer::new();

        // receive into response buffer
        let (_, src) = socket.recv_from(&mut response_buffer.buf).unwrap();

        let mut packet = DnsPacket::from_bytes(&mut response_buffer);

        // create the response packet
        let mut resp_packet = DnsPacket::empty();
        resp_packet.header.id = packet.header.id;
        resp_packet.header.recursion_desired = packet.header.recursion_desired;
        resp_packet.header.recursion_available = true;
        resp_packet.header.query_response = true;

        // if there is exactly one question
        if let Some(question) = packet.questions.pop() {
            println!("[*] Received query: {:?}", question);


            if let Ok(result) = DnsServer::lookup(&question.name, &question.qtype) {

                for rec  in result.answers {
                    println!("-> Answer: {:?}", rec);
                    resp_packet.answers.push(rec);
                }
                for rec  in result.authorities {
                    println!("-> Authority: {:?}", rec);
                    resp_packet.authorities.push(rec);
                }
                for rec  in result.resources {
                    println!("-> Resource: {:?}", rec);
                    resp_packet.resources.push(rec);
                }
                // if the lookup was successfull

                // push the original question
                resp_packet.questions.push(question);

            } else {
                resp_packet.header.result_code = ResultCode::SERVFAIL;
            }


        } else {
            resp_packet.header.result_code = ResultCode::FORMERR;
        }

        // write  the packet to a buffer and send via socket
        let mut res_buffer = BytePacketBuffer::new();
        resp_packet.write(&mut res_buffer).unwrap();

        socket.send_to(res_buffer.get_range(0, res_buffer.pos()).unwrap(), src).unwrap();


    }

    fn lookup(qname: &str, qtype: &QueryType) -> Result<DnsPacket, &'static str> {

        // use googles DNS mainly as primary DNS
        let server = ("8.8.8.8", 53);

        let mut pkt = DnsPacket::new(*qtype, qname);

        let mut req_buffer = BytePacketBuffer::new();
        pkt.write(&mut req_buffer).unwrap();

        let socket = UdpSocket::bind(("0.0.0.0", 43210)).unwrap();
        socket.send_to(&req_buffer.buf[0..req_buffer.pos], server).unwrap();

        let mut rec_buffer = BytePacketBuffer::new();
        socket.recv_from(&mut rec_buffer.buf).unwrap();

        let response = DnsPacket::from_bytes(&mut rec_buffer);

        return Ok(response);
    }
}


pub struct DnsTunnel {
    server: String,
    domain: String,
    port: u16,
    server_buffer: String
}

impl<'a> DnsTunnel {

    /// Craft a new instance of a DNS Tunnel
    /// ip: IP of the DNS tunnel server
    /// port: port of the DNS tunnel server
    pub fn new(domain: String, server: String, port: u16) -> DnsTunnel {
        DnsTunnel { server, port, domain, server_buffer: String::new() }
    }

    /// keeps open a listener to transfer data via DNS tunneling
    pub fn new_server(domain: String, port: u16) -> DnsTunnel {
        DnsTunnel { server: String::new(), port, domain, server_buffer: String::new() }
    }

    pub fn connect(&self) -> Result<(), &str> {
        if !self.check_connection() {
            println!("[-] DNS Tunnel not reachable");
            return Err("");
        }
        println!("[+] Successfully connected to DNS Tunnel");
        Ok(())
    }

    /// A function that can be called from the server side and lets the server send arbitrary
    /// strings
    pub fn server_send(&mut self, data: String) {
        self.server_buffer.push_str(&data);
    }

    pub fn send_string(&self, data: String) -> Result<(), &str> {

        if !self.init_send() {
            return Err("Failed to initialize data transfer");
        }

        // At this point, split the data in chunks of XX bytes and send each
        // chunk to the server until all the data is finished.
        // Close the connection by sending an `END` packet.
        if data.len() > 63 {
            // split the data in chunks of 63 bytes
            let chunks = data.as_bytes()
                .chunks(63)
                .map(str::from_utf8)
                .collect::<Result<Vec<&str>, _>>()
                .unwrap();

            for chunk in chunks {
                // Send the actual DATA after the tunnel is established
                println!("[DBG]: Sending Chunk: {}", chunk);
                // send the chunk
                self.send_chunk(chunk);
            }
        } else {
            self.send_chunk(&data);
        }


        // Close the data transfer by sending a `END` message
        if !self.close_send() {
            return Err("Failed to finish data transfer");
        }

        Ok(())

    }

    pub fn serve(&mut self) {

        let mut full_data: String = String::new();

        println!("[+] Started DNS tunnel, ready to receive data...");
        let socket = UdpSocket::bind(("0.0.0.0", self.port)).unwrap();

        // starting main loop
        let mut poll_status = PollStatus::NONE;

        let mut message = String::new();
        // UFF
        loop {
            let mut response_buffer = BytePacketBuffer::new();

            // receive into response buffer
            let (_, src) = socket.recv_from(&mut response_buffer.buf).unwrap();

            // decode the response as DNS packet
            let mut packet = DnsPacket::from_bytes(&mut response_buffer);

            // if there is at least one question of the type TXT
            if packet.questions.len() >= 1 && packet.questions[0].qtype == QueryType::TXT {
                let quest = &packet.questions[0];


                // Exfiltrate the Data
                let data = &packet.questions[0].name;
                let mut data: Vec<&str> = data.split(".").collect();
                let data = data[0];

                // Decode/Decrypt the data
                let data = b64decode(data);
                let data = String::from_utf8_lossy(&xor(&data, 0xF5)).to_string();


                // Return DNS Packet
                let mut return_packet = DnsPacket::empty();

                match &data[..] {
                    "INIT" => {
                        return_packet = DnsTunnel::get_txt_response(&packet, "ACK".to_string());
                    },
                    "START" => {
                        return_packet = DnsTunnel::get_txt_response(&packet, "ACK".to_string());
                    },
                    "END" => {
                        println!("Full Message {}", message);
                        message = String::new();
                        return_packet = DnsTunnel::get_txt_response(&packet, "ACK".to_string());
                    },
                    "POLL" => {
                        match poll_status {
                            PollStatus::NONE => {
                                // Enter Polling Mode
                                println!("[*] Polling Recieved");
                                return_packet = DnsTunnel::get_txt_response(&packet, "START".to_string());
                                poll_status = PollStatus::STARTED;
                            },
                            PollStatus::STARTED => {

                                let mut chunk = String::new();
                                if self.server_buffer.len() >= 63 {
                                    chunk = self.server_buffer[..63].to_string();
                                    self.server_buffer = self.server_buffer[63..].to_string();
                                } else {
                                    chunk.push_str(&self.server_buffer);
                                    self.server_buffer = String::new();
                                }
                                return_packet = DnsTunnel::get_txt_response(&packet, chunk);

                                if self.server_buffer.len() == 0 {
                                    poll_status = PollStatus::FINISHED;
                                }
                            },
                            PollStatus::FINISHED => {
                                // Send STOP message
                                return_packet = DnsTunnel::get_txt_response(&packet, "END".to_string());
                                poll_status = PollStatus::NONE;
                            }
                        }
                    },
                    _ => {
                        message.push_str(&data);
                    }

                }
                // return the response Packet
                let res_buffer = DnsTunnel::get_buffer(&mut return_packet);
                socket.send_to(res_buffer.get_range(0, res_buffer.pos()).unwrap(), src).unwrap();
            }
        }
    }



    /// this is not only checking for an `ACK` as answer, but gets the returned
    /// data
    fn send_recv(&self, message: &str) -> Result<String, ()> {

        // Take care of the encoding
        let bytes = xor(message.as_bytes(), 0xF5);
        let message = &b64encode(bytes);

        // Craft new request with `INIT`
        let mut pkt = self.get_txt_request(&self.domain, message.to_string());
        let resp = pkt.send((&self.server, self.port));

        // Check the answers
        if resp.answers.len() < 1 {
            return Err(());
        }

        let record = match &resp.answers[0] {
            DnsRecord::TXT {
                ref domain,
                len,
                ref txt,
                ttl
            } => {
                return Ok(txt.to_string())
            }
            _ => {
                    return Err(());
            }
        };
        return Err(());

    }
    /// Check if the remote server answers a simple `TUN` message with the
    /// answer `ACK`, this means that the server is available and can
    /// send back data.
    fn send_chunk(&self, message: &str) -> bool {


        // Take care of the encoding
        let bytes = xor(message.as_bytes(), 0xF5);
        let message = &b64encode(bytes);

        // Craft new request with `INIT`
        let mut pkt = self.get_txt_request(&self.domain, message.to_string());
        let resp = pkt.send((&self.server, self.port));

        // Check the answers
        if resp.answers.len() < 1 {
            return false;
        }

        let record = match &resp.answers[0] {
            DnsRecord::TXT {
                ref domain,
                len,
                ref txt,
                ttl
            } => {
                if txt == "ACK" {
                    return true;
                }
            }
            _ => {
                    return false;
            }
        };
        return false;
    }

    fn check_connection(&self) -> bool {
        self.send_chunk("INIT")
    }

    // Init the data transfer by sending a `START` message
    fn init_send(&self) -> bool {
        self.send_chunk("START")
    }
    // Closea the data transfer by sending an `END` message
    fn close_send(&self) -> bool {
        self.send_chunk("END")
    }

    pub fn receive_msg(&self) -> Option<String> {

        let mut wait_for_data = false;
        let mut return_data = String::new();

        loop {
            if let Ok(data) = self.send_recv("POLL") {
                match &data[..] {
                    "START" => {
                        println!("Start");
                        wait_for_data = true;
                    },
                    "END" => {
                        wait_for_data = false;
                        return Some(return_data);
                    },
                    "NONE" => {
                        return None;
                    },
                    _ => {
                        return_data.push_str(&data);
                    }

                }
            }
        }
        return None;
    }

    fn get_txt_response(ref_pkt: &DnsPacket, data: String) -> DnsPacket {
        let mut packet = DnsPacket::empty();
        packet.header.id = ref_pkt.header.id;
        packet.header.recursion_desired = ref_pkt.header.recursion_desired;
        packet.header.recursion_available = true;
        packet.header.query_response = true;
        packet.header.result_code = ResultCode::NOERROR;
        let domain = &ref_pkt.questions[0].name;

        let answer = DnsRecord::TXT {
            domain: domain.to_string(),
            len: data.len() as u8,
            ttl: 7766,
            txt: data
        };
        packet.answers.push(answer);
        return packet;

    }
    fn get_txt_request(&self, domain: &str, data: String) -> DnsPacket {
        let mut packet = DnsPacket::empty();
        let full_domain = format!("{}.{}", data, domain);
        let question = DnsQuestion{ name: full_domain, qtype: QueryType::TXT};
        packet.questions.push(question);

        packet
    }

    fn get_buffer(pkt: &mut DnsPacket) -> BytePacketBuffer {
        let mut res_buffer = BytePacketBuffer::new();
        pkt.write(&mut res_buffer).unwrap();
        return res_buffer;
    }

    /// Function splits  String ref into many chunks and returns a
    /// vector of those chunks
    fn split_data(&self, data: &'a String) -> Vec<&'a str> {

        if data.len() > 63 {
            // split the data in chunks of 63 bytes
            let chunks = data.as_bytes()
                .chunks(63)
                .map(str::from_utf8)
                .collect::<Result<Vec<&str>, _>>()
                .unwrap();
            return chunks;
        }
        let mut ret = Vec::new();
        ret.push(&data[..]);
        return ret;
    }
}


enum PollStatus {
    NONE,
    STARTED,
    FINISHED,
}