use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
use chrono::{Utc, TimeZone};
use data_encoding::{BASE32HEX, BASE64};
use rand::Rng;
use std::{collections::HashMap, cmp::max};
use std::io::{Cursor, Read, Seek, SeekFrom};
use std::fmt::{self, Display};
use strum_macros::EnumString;

pub mod error;
use error::{ParseError, EncodeError};

/// Represents a DNS OpCode.
/// See [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-5) for
/// further information.
#[derive(PartialEq, Copy, Clone, Debug)]
pub enum DnsOpcode {
    QUERY,
    IQUERY,
    STATUS,
    NOTIFY,
    UPDATE,
    DSO
}

/// Represents a DNS RCODE, including those introduced by EDNS.
/// See [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-6) for
/// further information.
#[derive(PartialEq, Copy, Clone, Debug)]
pub enum DnsRcode {
    NOERROR,
    FORMERR,
    SERVFAIL,
    NXDOMAIN,
    NOTIMP,
    REFUSED,
    YXDOMAIN,
    YXRRSET,
    NXRRSET,
    NOTAUTH,
    NOTZONE,
    DSOTYPENI,
    BADVERSBADSIG,
    BADKEY,
    BADTIME,
    BADMODE,
    BADNAME,
    BADALG,
    BADTRUNC,
    BADCOOKIE
}

/// Represents a DNS TYPE.
/// This enum is non-exhaustive, see [here](https://en.wikipedia.org/wiki/List_of_DNS_record_types) for a
/// more comprehensive overview as well as explanations and links to the respective defining RFCs.
#[derive(PartialEq, Copy, Clone, EnumString, Debug)]
pub enum DnsType {
    A,
    NS,
    CNAME,
    SOA,
    PTR,
    HINFO,
    MX,
    TXT,
    RP,
    // TODO: SIG (24)
    KEY,
    AAAA,
    LOC,
    SRV,
    NAPTR,
    CERT,
    DNAME,
    OPT,
    DS,
    SSHFP,
    // TODO: IPSECKEY (45)
    RRSIG,
    NSEC,
    DNSKEY,
    // TODO: DHCID (49)
    NSEC3,
    NSEC3PARAM,
    TLSA,
    // TODO: SMIMEA (53)
    // TODO: HIP (55)
    // TODO: CDNSKEY (60)
    OPENPGPKEY,
    // TODO: HTTPS (65)
    // TODO: TKEY (249)
    // TODO: TSIG (250)
    CAA
    // TODO: TA (32768)
    // TODO: DLV (32769)
}

/// Represents a DNS CLASS.
/// See [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) for further information.
/// Other classes than IN and ANY are included only for completeness and historical reasons.
#[derive(PartialEq, Copy, Clone, Debug)]
pub enum DnsClass {
    IN,
    CH,
    HS,
    NONE,
    ANY
}

/// Represents the flags of a `DnsHeader`.
pub struct DnsFlags {
    flags: u16
}

/// Represents a DNS header.
/// `rcode` is an `Option<DnsRcode>` because valid DNS queries do not have a RCODE.
/// See [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12) for
/// a list of available flags and their respective defining RFCs. The general format of a header is
/// defined in [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf).
pub struct DnsHeader {
    msg_id: u16,  // supplied by questioner and reflected back unchanged by responder
    qr: bool,     // false for query, true for response
    opcode: DnsOpcode,
    flags: DnsFlags,
    rcode: Option<DnsRcode>,
    qdcount: u16, // number of questions
    ancount: u16, // number of resource records
    nscount: u16, // number of name server resource records
    arcount: u16  // number of additional resource records
}

/// Represents a DNS question, i.e. an entry in the question section of a DNS message.
/// See [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) for further information.
pub struct DnsQuestion {
    qname: String,
    qtype: DnsType,
    qclass: DnsClass
}

/// Represents a DNS record, i. e. an entry in the answer, authority or additional section of a DNS
/// message.
/// See [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) for further information.
#[derive(Clone)]
pub enum DnsRecord {
    OPT {
        name: String,
        atype: DnsType,           // called `type` in RFC 1035 but `atype` here to avoid a conflict
                                  // with the Rust keyword
        payload_size: u16,
        rcode: DnsRcode,          // Always is the correct RCODE, i. e. the lower
                                  // four bits from the header are included
        edns_version: u8,
        flags: HashMap<&'static str, bool>,
                                  // rdlength omitted as rdata knows its own length
        rdata: Vec<u8>,           // needed for encoding
        parsed_rdata: Vec<String> // see DnsRecord::parse_rdata() for more information
    },
    NONOPT {
        name: String,
        atype: DnsType,
        class: DnsClass,
        ttl: u32,
                                  // rdlength omitted as rdata knows its own length
        rdata: Vec<u8>,           // needed for encoding
        parsed_rdata: Vec<String> // see DnsRecord::parse_rdata() for more information
    }
}

/// Represents a DNS message.
/// See [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) for further information.
pub struct DnsMessage {
    pub header: DnsHeader,
    questions: Vec<DnsQuestion>,
    answers: Vec<DnsRecord>,
    authoritative_answers: Vec<DnsRecord>,
    additional_answers: Vec<DnsRecord>
}

impl DnsOpcode {
    /// Encodes a DnsOpcode as a byte.
    pub fn encode(&self) -> u8 {
        match self {
            DnsOpcode::QUERY  => 0,
            DnsOpcode::IQUERY => 1,
            DnsOpcode::STATUS => 2,
            DnsOpcode::NOTIFY => 4,
            DnsOpcode::UPDATE => 5,
            DnsOpcode::DSO    => 6
        }
    }

    /// Parses an encoded DnsOpcode from a byte.
    /// Returns an error if the given byte does not represent a valid DNS OpCode.
    pub fn parse(val: u8) -> Result<DnsOpcode, ParseError> {
        Ok(match val {
            0 => DnsOpcode::QUERY,
            1 => DnsOpcode::IQUERY,
            2 => DnsOpcode::STATUS,
            4 => DnsOpcode::NOTIFY,
            5 => DnsOpcode::UPDATE,
            6 => DnsOpcode::DSO,
            x => return Err(ParseError::InvalidOpcode(x))
        })
    }
}

impl Display for DnsOpcode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl DnsRcode {
    /// Encodes a DnsRcode as a byte (actually only the lower four bits are used).
    /// Note that for RCODEs BADVERSBADSIG and following only the lower four bits are encoded;
    /// the upper eight bits need to be encoded in an OPT record in the additional section of the
    /// DNS message.
    pub fn encode(&self) -> u8 {
        match self {
            DnsRcode::NOERROR       => 0,
            DnsRcode::FORMERR       => 1,
            DnsRcode::SERVFAIL      => 2,
            DnsRcode::NXDOMAIN      => 3,
            DnsRcode::NOTIMP        => 4,
            DnsRcode::REFUSED       => 5,
            DnsRcode::YXDOMAIN      => 6,
            DnsRcode::YXRRSET       => 7,
            DnsRcode::NXRRSET       => 8,
            DnsRcode::NOTAUTH       => 9,
            DnsRcode::NOTZONE       => 10,
            DnsRcode::DSOTYPENI     => 11,
            DnsRcode::BADVERSBADSIG => 16 & 0b1111,
            DnsRcode::BADKEY        => 17 & 0b1111,
            DnsRcode::BADTIME       => 18 & 0b1111,
            DnsRcode::BADMODE       => 19 & 0b1111,
            DnsRcode::BADNAME       => 20 & 0b1111,
            DnsRcode::BADALG        => 21 & 0b1111,
            DnsRcode::BADTRUNC      => 22 & 0b1111,
            DnsRcode::BADCOOKIE     => 23 & 0b1111
        }
    }

    /// Parses an encoded DnsRcode from a twelve bit value. If EDNS is used, the upper eight bits
    /// are stored in the OPT entry of the additional section and the lower four bits are stored in
    /// the DnsHeader.
    /// Returns an error if the given value does not represent a valid DNS RCODE.
    pub fn parse(val: u16) -> Result<DnsRcode, ParseError> {
        Ok(match val {
            0  => DnsRcode::NOERROR,
            1  => DnsRcode::FORMERR,
            2  => DnsRcode::SERVFAIL,
            3  => DnsRcode::NXDOMAIN,
            4  => DnsRcode::NOTIMP,
            5  => DnsRcode::REFUSED,
            6  => DnsRcode::YXDOMAIN,
            7  => DnsRcode::YXRRSET,
            8  => DnsRcode::NXRRSET,
            9  => DnsRcode::NOTAUTH,
            10 => DnsRcode::NOTZONE,
            11 => DnsRcode::DSOTYPENI,
            16 => DnsRcode::BADVERSBADSIG,
            17 => DnsRcode::BADKEY,
            18 => DnsRcode::BADTIME,
            19 => DnsRcode::BADMODE,
            20 => DnsRcode::BADNAME,
            21 => DnsRcode::BADALG,
            22 => DnsRcode::BADTRUNC,
            23 => DnsRcode::BADCOOKIE,
            x => return Err(ParseError::InvalidRcode(x))
        })
    }
}

impl Display for DnsRcode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl DnsType {
    /// Encodes a DnsType as a two-byte value.
    pub fn encode(&self) -> u16 {
        match self {
            DnsType::A          => 1,
            DnsType::NS         => 2,
            DnsType::CNAME      => 5,
            DnsType::SOA        => 6,
            DnsType::PTR        => 12,
            DnsType::HINFO      => 13,
            DnsType::MX         => 15,
            DnsType::TXT        => 16,
            DnsType::RP         => 17,
            DnsType::KEY        => 25,
            DnsType::AAAA       => 28,
            DnsType::LOC        => 29,
            DnsType::SRV        => 33,
            DnsType::NAPTR      => 35,
            DnsType::CERT       => 37,
            DnsType::DNAME      => 39,
            DnsType::OPT        => 41,
            DnsType::DS         => 43,
            DnsType::SSHFP      => 44,
            DnsType::RRSIG      => 46,
            DnsType::NSEC       => 47,
            DnsType::DNSKEY     => 48,
            DnsType::NSEC3      => 50,
            DnsType::NSEC3PARAM => 51,
            DnsType::TLSA       => 52,
            DnsType::OPENPGPKEY => 61,
            DnsType::CAA        => 257
        }
    }

    /// Parses an encoded DnsType from a two-byte value.
    /// Returns an error if the given value does not represent a valid DNS TYPE or the represented
    /// TYPE has not been implemented.
    pub fn parse(val: u16) -> Result<DnsType, ParseError> {
        Ok(match val {
            1   => DnsType::A,
            2   => DnsType::NS,
            5   => DnsType::CNAME,
            6   => DnsType::SOA,
            12  => DnsType::PTR,
            13  => DnsType::HINFO,
            15  => DnsType::MX,
            16  => DnsType::TXT,
            17  => DnsType::RP,
            25  => DnsType::KEY,
            28  => DnsType::AAAA,
            29  => DnsType::LOC,
            33  => DnsType::SRV,
            35  => DnsType::NAPTR,
            37  => DnsType::CERT,
            39  => DnsType::DNAME,
            41  => DnsType::OPT,
            43  => DnsType::DS,
            44  => DnsType::SSHFP,
            46  => DnsType::RRSIG,
            47  => DnsType::NSEC,
            48  => DnsType::DNSKEY,
            50  => DnsType::NSEC3,
            51  => DnsType::NSEC3PARAM,
            52  => DnsType::TLSA,
            61  => DnsType::OPENPGPKEY,
            257 => DnsType::CAA,
            x => return Err(ParseError::InvalidType(x))
        })
    }

    /// Returns the schema used for parsing in `DnsRecord::parse_rdata()`.
    pub fn rdata_schema(&self) -> &str {
        match self {
            DnsType::A            => "ip4",
            DnsType::NS
                | DnsType::CNAME
                | DnsType::DNAME
                | DnsType::PTR    => "qname",
            DnsType::SOA          => "qname qname u32 u32 u32 u32 u32",
            DnsType::HINFO        => "string string",
            DnsType::MX           => "u16 qname",
            DnsType::TXT          => "text",
            DnsType::RP           => "qname qname",
            DnsType::KEY
                | DnsType::DNSKEY => "u16 u8 u8 base64",
            DnsType::AAAA         => "ip6",
            DnsType::LOC          => "u8 u8 u8 u8 u32 u32 u32",
            DnsType::SRV          => "u16 u16 u16 qname",
            DnsType::NAPTR        => "u16 u16 string string string qname",
            DnsType::CERT         => "u16 u16 u8 base64",
            DnsType::OPT          => "options",
            DnsType::DS           => "u16 u8 u8 hex",
            DnsType::SSHFP        => "u8 u8 hex",
            DnsType::RRSIG        => "qtype u8 u8 u32 time time u16 qname base64",
            DnsType::NSEC         => "qname types",
            DnsType::NSEC3        => "u8 u8 u16 salt hash types",
            DnsType::NSEC3PARAM   => "u8 u8 u16 salt",
            DnsType::TLSA         => "u8 u8 u8 hex",
            DnsType::OPENPGPKEY   => "base64",
            DnsType::CAA          => "u8 property"
        }
    }
}

impl Display for DnsType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl DnsClass {
    /// Encodes a DnsClass as a two-byte value.
    pub fn encode(&self) -> u16 {
        match self {
            DnsClass::IN   => 1,
            DnsClass::CH   => 3,
            DnsClass::HS   => 4,
            DnsClass::NONE => 254,
            DnsClass::ANY  => 255
        }
    }

    /// Parses an encoded DnsClass from a two-byte value.
    /// Returns an error if the given value does not represent a valid DNS CLASS.
    pub fn parse(val: u16) -> Result<DnsClass, ParseError> {
        Ok(match val {
            1   => DnsClass::IN,
            3   => DnsClass::CH,
            4   => DnsClass::HS,
            254 => DnsClass::NONE,
            255 => DnsClass::ANY,
            x => return Err(ParseError::InvalidClass(x))
        })
    }
}

impl Display for DnsClass {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl DnsFlags {
    /// Creates a `DnsFlags` struct from booleans. The meaning of the flags is as follows:
    /// * `aa`: authoritative answer (valid in responses only)
    /// * `tc`: truncated (set on all truncated messages except last one)
    /// * `rd`: recursion desired (copied in answer if supported and accepted)
    /// * `ra`: valid in responses, indicating recursive query support in the name server
    /// * `ad`: if set in query: indicates interest in the ad bit of the upcoming response;
    ///     if set in response: indicates that the resolver side considers all RRsets in
    ///     the Answer section and relevant negative response RRs in the Authority section
    ///     to be authentic.
    /// * `cd`: disable signature validation in a security-aware name server's processing of a particular query
    pub fn new(aa: bool, tc: bool, rd: bool, ra: bool, ad: bool, cd: bool) -> Self {
        let aa = if aa {1} else {0};
        let tc = if tc {1} else {0};
        let rd = if rd {1} else {0};
        let ra = if ra {1} else {0};
        let ad = if ad {1} else {0};
        let cd = if cd {1} else {0};
        DnsFlags {
            flags: (aa << 10) + (tc << 9) + (rd << 8) + (ra << 7) + (ad << 5) + (cd << 4)
        }
    }

    /// Creates a `DnsFlags` struct from bitflags as they would appear in the second 16-octet line of a `DnsHeader`.
    pub fn from_flags(flags: u16) -> Self {
        DnsFlags { flags }
    }

    /// Returns a u16 representing bitflags as they would appear in the second 16-octet line of a `DnsHeader`.
    pub fn as_flags(&self) -> u16 {
        self.flags
    }

    /// Whether the aa flag is set.
    pub fn aa(&self) -> bool { (self.flags & (1 << 10)) != 0 }

    /// Whether the tc flag is set.
    pub fn tc(&self) -> bool { (self.flags & (1 << 9)) != 0 }

    /// Whether the rd flag is set.
    pub fn rd(&self) -> bool { (self.flags & (1 << 8)) != 0 }

    /// Whether the ra flag is set.
    pub fn ra(&self) -> bool { (self.flags & (1 << 7)) != 0 }

    /// Whether the ad flag is set.
    pub fn ad(&self) -> bool { (self.flags & (1 << 5)) != 0 }

    /// Whether the cd flag is set.
    pub fn cd(&self) -> bool { (self.flags & (1 << 4)) != 0 }
}

impl DnsHeader {
    /// Creates a header for a DNS response message.
    /// See [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) and
    /// [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12) for
    /// information about the parameters. `qdcount`, `ancount`, `nscount` and `arcount` are grouped in that
    /// order in the `counts` parameter.
    pub fn new_response_header(msg_id: u16, opcode: DnsOpcode, flags: DnsFlags, rcode: DnsRcode, counts: [u16; 4]) -> Self {
        DnsHeader {
            msg_id,
            qr: true,
            opcode,
            flags,
            rcode: Some(rcode),
            qdcount: counts[0],
            ancount: counts[1],
            nscount: counts[2],
            arcount: counts[3]
        }
    }

    /// Creates a header for a DNS query message.
    /// If the query includes an OPT record, `edns` must be `true`. See
    /// [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) and
    /// [here](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12) for
    /// information about the other parameters.
    pub fn new_query_header(msg_id: u16, opcode: DnsOpcode, flags: DnsFlags,
                            edns: bool, qdcount: u16) -> Result<Self, EncodeError> {
        if flags.aa() || flags.ra() {
            Err(EncodeError::AaOrRaInQuery)
        } else {
            Ok(DnsHeader {
                msg_id,
                qr: false,
                opcode,
                flags,
                rcode: None,
                qdcount,
                ancount: 0,
                nscount: 0,
                arcount: if edns {1} else {0}
            })
        }
    }

    /// Encodes a DnsHeader as a series of bytes.
    /// Returns an error if a method defined in `byteorder::WriteBytesExt` returns an error.
    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
        let mut res = Vec::new();
        let qr = if self.qr {1u16} else {0u16};
        let opcode = self.opcode.encode() as u16;
        let rcode = match &self.rcode {
            Some(val) => val.encode() as u16,
            None => 0u16
        };

        let line_two = (qr << 15) + (opcode << 11) + self.flags.as_flags() + rcode;
        res.write_u16::<NetworkEndian>(self.msg_id)?;
        res.write_u16::<NetworkEndian>(line_two)?;
        res.write_u16::<NetworkEndian>(self.qdcount)?;
        res.write_u16::<NetworkEndian>(self.ancount)?;
        res.write_u16::<NetworkEndian>(self.nscount)?;
        res.write_u16::<NetworkEndian>(self.arcount)?;

        Ok(res)
    }

    /// Parses an encoded DnsHeader from a series of bytes.
    /// Returns an error if `DnsOpcode::parse()`, `DnsRcode::parse()` or a method defined in
    /// `byteorder::ReadBytesExt` return an error.
    pub fn parse(header: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
        let msg_id =  header.read_u16::<NetworkEndian>()?;
        let line_two = header.read_u16::<NetworkEndian>()?;
        let qr = (line_two & (1 << 15)) >> 15;
        let opcode = DnsOpcode::parse(((line_two & (0b1111 << 11)) >> 11) as u8)?;
        let flags = DnsFlags::from_flags(line_two & 0b0000011110110000);
        let rcode = DnsRcode::parse(line_two & 0b1111)?;

        Ok(DnsHeader {
            msg_id,
            qr: qr != 0,
            opcode,
            flags,
            rcode: if qr != 0 {Some(rcode)} else {None},
            qdcount: header.read_u16::<NetworkEndian>()?,
            ancount: header.read_u16::<NetworkEndian>()?,
            nscount: header.read_u16::<NetworkEndian>()?,
            arcount: header.read_u16::<NetworkEndian>()?
        })
    }

    /// Creates a string containing information (id, opcode, rcode if applicable, flags) about the header.
    pub fn info_str(&self) -> String {
        let mut s = String::new();
        if let Some(rcode) = self.rcode {
            s.push_str(format!("id: {}, opcode: {}, rcode: {}, flags: ", self.msg_id, self.opcode,
                               rcode).as_str());
        } else {
            s.push_str(format!("id: {}, opcode: {}, flags: ", self.msg_id, self.opcode).as_str());
        }
        if self.flags.aa() { s.push_str("aa ") }
        if self.flags.tc() { s.push_str("tc ") }
        if self.flags.rd() { s.push_str("rd ") }
        if self.flags.ra() { s.push_str("ra ") }
        if self.flags.ad() { s.push_str("ad ") }
        if self.flags.cd() { s.push_str("cd ") }
        s.remove(s.len() - 1); // remove last ' '
        s
    }
}

impl Display for DnsHeader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut s = String::new();
        if self.qr {
            s.push_str("DNS Response (");
        } else {
            s.push_str("DNS Query (");
        }
        s.push_str(&self.info_str());
        s.push(')');
        write!(f, "{}", s)
    }
}

impl DnsQuestion {
    /// Creates a DNS question.
    pub fn new(domain: &str, qtype: DnsType, qclass: DnsClass) -> Self {
        DnsQuestion {
            qname: domain.to_string(),
            qtype,
            qclass
        }
    }

    /// Encodes a DnsQuestion as a series of bytes.
    /// Returns an error if a method defined in `byteorder::WriteBytesExt` returns an error.
    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
        let mut question = DnsMessage::encode_qname(self.qname.as_str())?;
        question.write_u16::<NetworkEndian>(self.qtype.encode())?;
        question.write_u16::<NetworkEndian>(self.qclass.encode())?;
        Ok(question)
    }

    /// Parses an encoded DnsQuestion from a series of bytes.
    /// Returns an error if `DnsMessage::parse_qname()`, `DnsType::parse()`, `DnsClass::parse()`
    /// or a method defined in `byteorder::ReadBytesExt` return an error.
    pub fn parse(msg: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
        let qname = DnsMessage::parse_qname(msg)?;
        let qtype = DnsType::parse(msg.read_u16::<NetworkEndian>()?)?;
        let qclass = DnsClass::parse(msg.read_u16::<NetworkEndian>()?)?;

        Ok(DnsQuestion {
            qname,
            qtype,
            qclass
        })
    }

    /// Returns a string representing the record in the canonical format, with the
    /// owner padded to the given length.
    pub fn as_padded_string(&self, owner_len: usize) -> String {
        let mut res = String::new();
        let mut owner = self.qname.clone();
        while owner.len() < owner_len {
            owner.push(' ');
        }
        res.push_str(format!(
            "{}\t    \t{}\t{}", owner, self.qclass, self.qtype // four spaces to mimic ttl
        ).as_str());

        res
    }
}

impl Display for DnsQuestion {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "DNS Question for '{}' (type: {}, class: {})", self.qname, self.qtype, self.qclass)
    }
}

impl DnsRecord {
    // TODO: maybe implement support for some options?
    /// Creates a new OPT record.
    pub fn new_opt_record(payload_size: u16, do_flag: bool) -> Self {
        let mut flags = HashMap::new();
        flags.insert("do", do_flag);
        DnsRecord::OPT {
            name: "".to_string(),
            atype: DnsType::OPT,
            payload_size,
            rcode: DnsRcode::NOERROR,
            edns_version: 0,
            flags,
            rdata: vec![],
            parsed_rdata: Default::default()
        }
    }

    /// Encodes a DnsRecord as a series of bytes.
    /// Returns an error if a method defined in `byteorder::WriteBytesExt` returns an error.
    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
        match self {
            DnsRecord::NONOPT { name, atype, class, ttl, rdata, .. } => {
                let mut record =  DnsMessage::encode_qname(name.as_str())?;
                record.write_u16::<NetworkEndian>(atype.encode())?;
                record.write_u16::<NetworkEndian>(class.encode())?;
                record.write_u32::<NetworkEndian>(*ttl)?;
                record.write_u16::<NetworkEndian>(rdata.len() as u16)?;
                record.append(&mut rdata.clone());
                Ok(record)
            },
            DnsRecord::OPT { name, atype, payload_size, rcode, edns_version, flags,
                rdata, .. } => {
                let mut record =  DnsMessage::encode_qname(name.as_str())?;
                record.write_u16::<NetworkEndian>(atype.encode())?;
                record.write_u16::<NetworkEndian>(*payload_size)?;
                let rcode = (((rcode.encode() as u16) & 0b111111110000) >> 4) as u8;
                record.write_u8(rcode)?;
                record.write_u8(*edns_version)?;
                if flags.contains_key("do") && flags["do"] {
                    record.write_u16::<NetworkEndian>(1 << 15)?;
                } else {
                    record.write_u16::<NetworkEndian>(0)?;
                }
                record.write_u16::<NetworkEndian>(rdata.len() as u16)?;
                record.append(&mut rdata.clone());
                Ok(record)
            }
        }
    }

    /// Parses an encoded DnsRecord from a series of bytes.
    /// Returns an error if `DnsMessage::parse_qname()`, `DnsType::parse()`, `DnsClass::parse()`,
    /// `DnsRecord::parse_rdata()` or a method defined in `byteorder::ReadBytesExt` return an error
    /// or an OPT record has a name other than `""`.
    pub fn parse(msg: &mut Cursor<&[u8]>, rcode: DnsRcode) -> Result<Self, ParseError> {
        let name = DnsMessage::parse_qname(msg)?;
        let t = msg.read_u16::<NetworkEndian>()?;
        let atype = DnsType::parse(t)?;
        if atype == DnsType::OPT {
            return DnsRecord::parse_opt_record(msg, name, rcode);
        }
        let class = DnsClass::parse(msg.read_u16::<NetworkEndian>()?)?;
        let ttl = msg.read_u32::<NetworkEndian>()?;
        let rdlength = msg.read_u16::<NetworkEndian>()?;

        let mut rdata = vec![0; rdlength as usize];
        let pos_rdata_start = msg.position();
        msg.read_exact(&mut rdata)?;
        // reset position to the start of rdata for parse_rdata()
        msg.set_position(pos_rdata_start);
        let parsed_rdata = DnsRecord::parse_rdata(&atype, msg, rdlength)?;

        Ok(DnsRecord::NONOPT {
            name,
            atype,
            class,
            ttl,
            rdata,
            parsed_rdata
        })
    }

    /// Parses encoded rdata into a vector of strings (canonical format).
    /// `atype` is the type of the record containing the rdata. `msg` is the complete response
    /// message, which is needed for message compression. `rdlength` is the length of the RDATA in
    /// bytes.
    /// Returns an error if `DnsMessage::parse_qname()`, `DnsMessage::parse_string()`, `DnsType::parse()`,
    /// `DnsClass::parse()`, `DnsRecord::parse_rdata()`, `DnsRecord::parse_nsec_type_bitmap()`,
    /// `DnsRecord::interpret_nsec_type_bitmap()` or a method defined in `byteorder::ReadBytesExt`
    /// return an error or a TXT record contains an invalid byte.
    pub fn parse_rdata(atype: &DnsType, msg: &mut Cursor<&[u8]>, rdlength: u16) -> Result<Vec<String>, ParseError> {
        let mut res = Vec::new();
        let schema = atype.rdata_schema();
        let mut len_read = 0u16;
    
        for token in schema.split(' ') {
            let pos_before = msg.position();
            res.push(match token {
                "u8"     => { msg.read_u8()?.to_string() },
                "u16"    => { msg.read_u16::<NetworkEndian>()?.to_string() },
                "u32"    => { msg.read_u32::<NetworkEndian>()?.to_string() },
                "qname"  => { DnsMessage::parse_qname(msg)? },
                "string" => { format!("\"{}\"", DnsMessage::parse_string(msg)?) },
                "ip4"    => { format!("{}.{}.{}.{}", msg.read_u8()?, msg.read_u8()?, msg.read_u8()?, msg.read_u8()?) },
                "ip6"    => {
                    let mut addr = String::new();
                    for _ in 0..8 {
                        addr.push_str(format!("{:x}:", msg.read_u16::<NetworkEndian>()?).as_str());
                    }
                    addr.remove(addr.len() - 1); // remove last ':'
                    addr
                },
                "text" => {
                    let mut s = String::new();
                    let mut len = 0;
                    // according to RFC1035 it is possible that one TEXT entry holds multiple character strings
                    while len < rdlength - len_read {
                        let t = DnsMessage::parse_string(msg)?;
                        s.push_str(t.as_str());
                        len += (t.len() as u16) + 1; // also count the length byte before the actual string
                    }
                    s
                },
                "hex" => {
                    let mut hex = String::new();
                    for _i in 0..(rdlength - len_read) {
                        hex.push_str(format!("{:02x}", msg.read_u8()?).as_str());
                    }
                    hex
                },
                "qtype" => {
                    DnsType::parse(msg.read_u16::<NetworkEndian>()?)?.to_string()
                },
                "base64" => {
                    let mut data = vec![0; (rdlength - len_read) as usize];
                    msg.read_exact(&mut data)?;
                    BASE64.encode(&data)
                },
                // NSEC/NSEC3 type bitmap
                "types" => {
                    let bitmap = DnsRecord::parse_nsec_type_bitmap(msg, len_read, rdlength)?;
                    DnsRecord::interpret_nsec_type_bitmap(bitmap)?
                },
                // NSEC3/NSEC3PARAM salt (including length octet)
                "salt" => {
                    let salt_len = msg.read_u8()?;
                    if salt_len == 0 {
                        "-".to_string()
                    } else {
                        let mut salt = String::new();
                        for _ in 0..salt_len {
                            salt.push_str(format!("{:02x}", msg.read_u8()?).as_str());
                        }
                        salt
                    }
                },
                // NSEC3 hashed owner (including length octet)
                "hash" => {
                    let hash_len = msg.read_u8()? as usize;
                    let mut hash = vec![0; hash_len];
                    msg.read_exact(&mut hash)?;
                    BASE32HEX.encode(&hash)
                },
                // CAA Property, i. e. tag-value pair (inlcuding length octet)
                "property" => {
                    let tag_len = msg.read_u8()?;
                    let mut property = String::new();
                    for _i in 0..tag_len {
                        property.push(msg.read_u8()? as char);
                    }
                    property.push(' ');
                    for _i in 0..(rdlength - (tag_len as u16) - 1 - len_read) {
                        property.push(msg.read_u8()? as char);
                    }
                    property
                },
                // OPT options
                "options" => {
                    let mut len = 0;
                    let mut s = String::new();
                    while len < rdlength - len_read {
                        let option_code = msg.read_u16::<NetworkEndian>()?;
                        s.push_str(option_code.to_string().as_str());
                        s.push_str(": ");
                        let option_len = msg.read_u16::<NetworkEndian>()?;
                        for _ in 0..option_len {
                            s.push_str(format!("{:02x}", msg.read_u8()?).as_str());
                        }
                        s.push_str(", ");
                        len += option_len + 4;
                    }
                    // not all OPT records have rdata
                    if !s.is_empty() {
                        s.remove(s.len() - 1); // remove last ' '
                        s.remove(s.len() - 1); // remove last ','
                    }
                    s
                },
                // RRSIG signature expiration/inception
                "time" => {
                    let ts = msg.read_u32::<NetworkEndian>()?;
                    Utc.timestamp(ts as i64, 0).format("%Y%m%d%H%M%S").to_string()
                },
                x => {
                    panic!(format!("Invalid rdata schema key: {}", x))
                }
            });
            len_read += (msg.position() - pos_before) as u16;
        }
        Ok(res)
    }

    /// Takes a list of DNS TYPEs (e. g. the return value of `DnsRecord::parse_nsec_type_bitmap()`)
    /// and returns a string containing the representation of the listed TYPEs.
    /// Returns an error if `DnsType::parse()` returns an error (i. e. `types` contains an invalid
    /// or unimplemented type).
    pub fn interpret_nsec_type_bitmap(types: Vec<u16>) -> Result<String, ParseError> {
        let mut res = String::new();
        for t in types {
            let t = match DnsType::parse(t) {
                Ok(r) => r.to_string(),
                Err(_) => format!("TYPE{}", t) // RFC 3597, Section 5
            };
            res.push_str(t.as_str());
            res.push(' ');
        }
        // apparently sometimes an nsec(3) record includes no types
        if !res.is_empty() {
            res.remove(res.len() - 1); // remove last ' '
        }
        Ok(res)
    }

    /// Parses the type bitmap in the RDATA section of an NSEC or NSEC3 record.
    /// `len_read` is the count of the bytes already read from the rdata. `rdlength` is the
    /// total length of the rdata.
    /// Returns an error if a method defined in `byteorder::ReadBytesExt` returns an error.
    pub fn parse_nsec_type_bitmap(msg: &mut Cursor<&[u8]>, len_read: u16, rdlength: u16)
        -> Result<Vec<u16>, ParseError> {
        let mut len_read = len_read;
        let mut available_types = Vec::new();
        while len_read < rdlength {
            let window_number = msg.read_u8()?;
            let bitmap_len = msg.read_u8()?;
            for i in 0..bitmap_len {
                let byte = msg.read_u8()?;
                for j in 0..8 {
                    if (byte & (0b10000000 >> j)) != 0 {
                        let type_num = ((window_number as u16) << 8) + (i * 8 + j) as u16;
                        available_types.push(type_num);
                    }
                }
            }
            len_read += (2 + bitmap_len) as u16;
        }
        Ok(available_types)
    }

    /// Returns a string representing the record in the canonical format, with the
    /// owner padded to the given length.
    /// Panics when called on an OPT record.
    pub fn as_padded_string(&self, owner_len: usize) -> String {
        match self {
            DnsRecord::NONOPT { name, atype, class, ttl, parsed_rdata, .. } => {
                let mut res = String::new();
                let mut owner = name.clone();
                while owner.len() < owner_len {
                    owner.push(' ');
                }
                res.push_str(format!(
                    "{}\t{}\t{}\t{}\t\t{}", owner, ttl, class, atype,
                    DnsRecord::format_parsed_rdata(parsed_rdata)
                ).as_str());

                res
            },
            DnsRecord::OPT { .. } => {
                panic!()
            }
        }
    }

    /// Parses an encoded OPT record from a series of bytes.
    /// See DnsRecord::parse() for further information.
    fn parse_opt_record(msg: &mut Cursor<&[u8]>, name: String, rcode: DnsRcode)
                        -> Result<DnsRecord, ParseError> {
        if !name.eq("") {
            return Err(ParseError::InvalidOptName(name))
        }

        let payload_size = msg.read_u16::<NetworkEndian>()?;
        let ext_rcode = msg.read_u8()?;
        let rcode = match ext_rcode {
            0 => rcode,
            x => DnsRcode::parse(((x as u16) << 4) + (rcode.encode() as u16))?
        };
        let edns_version = msg.read_u8()?;
        let mut flags = HashMap::new();
        flags.insert(
            "do",
            msg.read_u16::<NetworkEndian>()? & (1 << 15) != 0
        );

        let rdlength = msg.read_u16::<NetworkEndian>()?;
        let mut rdata = vec![0; rdlength as usize];
        let pos_rdata_start = msg.position();
        msg.read_exact(&mut rdata)?;
        // reset position to the start of rdata for parse_rdata()
        msg.set_position(pos_rdata_start);
        let parsed_rdata = DnsRecord::parse_rdata(&DnsType::OPT, msg, rdlength)?;

        Ok(DnsRecord::OPT {
            name,
            atype: DnsType::OPT,
            payload_size,
            rcode,
            edns_version,
            flags,
            rdata,
            parsed_rdata
        })
    }

    /// Formats parsed rdata into a string.
    fn format_parsed_rdata(parsed_rdata: &[String]) -> String {
        let mut res = String::new();
        for val in parsed_rdata {
            res.push_str(format!("{} ", val).as_str());
        }
        if !parsed_rdata.is_empty() {
            res.remove(res.len() - 1); // remove last ' '
        }
        res
    }
}

impl Display for DnsRecord {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DnsRecord::NONOPT { name, atype, class, ttl, parsed_rdata, ..} => {
                write!(
                    f, "DNS Record for '{}' (type: {}, class: {}, ttl: {}, rdata: {})",
                    name, atype, class, ttl, DnsRecord::format_parsed_rdata(parsed_rdata)
                )
            },
            DnsRecord::OPT { payload_size, edns_version, flags,
                parsed_rdata, .. } => {
                let mut s = format!(
                    "DNS OPT Record (EDNS version: {}, payload size: {}, flags: ",
                   edns_version, payload_size
                );
                for (key, value) in flags {
                    if *value { s.push_str(key); s.push(' '); }
                }
                s.remove(s.len() - 1); // remove last ' '
                s.push_str(format!(
                    ", rdata: {})", DnsRecord::format_parsed_rdata(parsed_rdata)
                ).as_str());
                write!(f, "{}", s)
            }
        }
    }
}

impl DnsMessage {
    /// Creates a DNS query.
    /// `edns` dictates whether the query contains an OPT record for indication of EDNS support.
    /// `bufsize` is the payload size that gets sent in the OPT record. If `edns == false` the value
    /// of `bufsize` is ignored. `do_flag` indicates DNSSEC support. See
    /// [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf) and the documentation of `DnsHeader` for
    /// information about the remaining parameters.
    /// Returns an error if `do_flag && !edns == true`.
    pub fn new_query(domain: &str, qtype: DnsType, opcode: DnsOpcode, flags: DnsFlags,
                     edns: bool, do_flag: bool, bufsize: u16) -> Result<Self, EncodeError> {
        if do_flag && !edns {
            return Err(EncodeError::DoSetButNoEdns);
        }

        if flags.aa() || flags.ra() {
            return Err(EncodeError::AaOrRaInQuery);
        }

        let msg_id = rand::thread_rng().gen_range(0..(1u32 << 16)) as u16;

        let mut additional_answers = Vec::new();
        if edns {
            additional_answers.push(DnsRecord::new_opt_record(bufsize, do_flag));
        }

        Ok(DnsMessage {
            header: DnsHeader::new_query_header(msg_id, opcode, flags, edns, 1)?,
            questions: vec![DnsQuestion::new(domain, qtype, DnsClass::IN)],
            answers: Vec::new(),
            authoritative_answers: Vec::new(),
            additional_answers
        })
    }

    /// Creates a DNS response.
    /// See the documentation of `DnsHeader` for information about the parameters.
    /// `answers`, `authoritative_answers`, `additional_answers` are grouped in that order in the `records` parameter.
    pub fn new_response(msg_id: u16, opcode: DnsOpcode, flags: DnsFlags, rcode: DnsRcode, questions: Vec<DnsQuestion>,
                        records: [Vec<DnsRecord>; 3]) -> Self {
        DnsMessage {
            header: DnsHeader::new_response_header(
                msg_id, opcode, flags, rcode,
                [questions.len() as u16, records[0].len() as u16, records[1].len() as u16, records[2].len() as u16]
            ),
            questions,
            answers: records[0].clone(),
            authoritative_answers: records[1].clone(),
            additional_answers: records[2].clone()
        }
    }

    /// Encodes a DnsMessage as a series of bytes.
    /// Returns an error if `DnsHeader::encode()`, `DnsQuestion::encode()` or `DnsRecord::encode()`
    /// return an error.
    pub fn encode(&self) -> Result<Vec<u8>, EncodeError> {
        let mut res = self.header.encode()?;
        for question in &self.questions {
            res.append(&mut question.encode()?);
        }
        for record in &self.answers {
            res.append(&mut record.encode()?);
        }
        for record in &self.authoritative_answers {
            res.append(&mut record.encode()?);
        }
        for record in &self.additional_answers {
            res.append(&mut record.encode()?);
        }

        Ok(res)
    }

    /// Parses an encoded DnsMessage from a series of bytes.
    /// Returns an error if `DnsHeader::parse()`, `DnsQuestion::parse()` or `DnsRecord::parse()`
    /// return an error, the header/EDNS header contain an RCODE other than `DnsRcode::NoError` or a
    /// truncated message is received.
    pub fn parse(msg: &mut Cursor<&[u8]>) -> Result<Self, ParseError> {
        let mut header = DnsHeader::parse(msg)?;

        if header.flags.tc() {
            return Err(ParseError::TruncatedMessage);
        }

        let qdcount = header.qdcount;
        let ancount = header.ancount;
        let nscount = header.nscount;
        let arcount = header.arcount;
        let questions = DnsMessage::parse_questions(msg, qdcount)?;
        let mut answers = Vec::new();
        let mut authoritative_answers = Vec::new();
        let mut additional_answers = Vec::new();
        if ancount > 0 {
            answers = DnsMessage::parse_records(
                msg, ancount, header.rcode.unwrap()
            )?;
        }
        if nscount > 0 {
            authoritative_answers = DnsMessage::parse_records(
                msg, nscount, header.rcode.unwrap()
            )?;
        }
        if arcount > 0 {
            additional_answers = DnsMessage::parse_records(
                msg, arcount, header.rcode.unwrap()
            )?;
        }

        for answer in &additional_answers {
            if let DnsRecord::OPT { rcode, .. } = answer {
                header.rcode = Some(*rcode);
            }
        }

        Ok(DnsMessage {
            header,
            questions,
            answers,
            authoritative_answers,
            additional_answers
        })
    }

    /// Encodes a domain name as a DNS QNAME. Does not use message compression.
    /// Returns an error if the domain name is longer than 255 characters, a label
    /// (part between two dots) in the domain name is longer than 63 characters or a method defined
    /// in `byteorder::WriteBytesExt` returns an error.
    pub fn encode_qname(domain: &str) -> Result<Vec<u8>, EncodeError> {
        if domain.bytes().len() > 255 {
            return Err(EncodeError::DomainTooLong(domain.bytes().len())
            );
        }

        if domain.eq("") { // special case for the root label
            return Ok(vec![0]);
        }

        let mut res = Vec::new();
        for label in domain.split('.') {
            if label.bytes().len() > 63 {
                return Err(EncodeError::LabelTooLong(label.bytes().len())
                );
            }
            res.write_u8(label.len() as u8)?;
            label.bytes().for_each(|b| res.push(b));
        }
        res.write_u8(0)?;

        Ok(res)
    }

    /// Parses a character string as defined by [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf).
    /// Returns an error if a method defined in `byteorder::ReadBytesExt` returns an error.
    pub fn parse_string(msg: &mut Cursor<&[u8]>) -> Result<String, ParseError> {
        let length = msg.read_u8()?;
        let mut res = String::new();
        for _i in 0..length {
            res.push(msg.read_u8()? as char);
        }
        Ok(res)
    }

    /// Parses a DNS QNAME as defined by [RFC 1035](https://tools.ietf.org/pdf/rfc1035.pdf),
    /// returning it in common url format. As required by the specification, method compression
    /// is supported.
    /// Returns an error if a method defined in `byteorder::ReadBytesExt` returns an error or an
    /// extended or invalid label type is encountered.
    pub fn parse_qname(msg: &mut Cursor<&[u8]>) -> Result<String, ParseError> {
        let mut domain = String::new();
        let mut c = msg.read_u8()?; // length of next label

        while c != 0 {
            if (c & 0b11000000) != 0 { // after this comes a pointer for message compression
                c &= 0b00111111; // erase upper two bits of c for offset calculation
                let offset = ((c as u16) << 8) + (msg.read_u8()? as u16);
                // save position after pointer
                let pos_after_pointer = msg.position() as i64;
                msg.seek(SeekFrom::Start(offset as u64))?;
                // recursion is the easiest way to handle recursive message compression
                // (i've seen that being used... looking at you, a.gtld-servers.net)
                domain.push_str(DnsMessage::parse_qname(msg)?.as_str());

                // move cursor to byte after pointer
                msg.seek(SeekFrom::Start(pos_after_pointer as u64))?;
                return Ok(domain);
            } else if (c & 0b01000000) != 0 || (c & 0b10000000) != 0 {
                return Err(ParseError::InvalidLabelType(c));
            }
            for _i in 0..c {
                domain.push(msg.read_u8()? as char);
            }
            domain.push('.');
            c = msg.read_u8()?;
        }

        Ok(domain)
    }

    /// Parses the question section of a DNS message.
    fn parse_questions(msg: &mut Cursor<&[u8]>, qdcount: u16) -> Result<Vec<DnsQuestion>, ParseError> {
        let mut questions = Vec::with_capacity(qdcount as usize);
        for _i in 0..qdcount {
            questions.push(DnsQuestion::parse(msg)?);
        }

        Ok(questions)
    }

    /// Parses an answer section (i. e. answer, authoritative or additional) of a DNS message.
    fn parse_records(msg: &mut Cursor<&[u8]>, ancount: u16, rcode: DnsRcode) -> Result<Vec<DnsRecord>, ParseError> {
        let mut answers = Vec::with_capacity(ancount as usize);
        for _i in 0..ancount {
            answers.push(DnsRecord::parse(msg, rcode)?);
        }

        Ok(answers)
    }
}

impl Display for DnsMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut res = String::new();

        let mut additional_answers = self.additional_answers.clone();
        let mut opt = None;
        let mut opt_index = 0;

        let mut max_owner_len = 0;
        let answers = [&self.answers, &self.authoritative_answers, &self.additional_answers];
        let answers_iter = answers
            .iter()
            .flat_map(|a| a.iter());
        for (i, answer) in answers_iter.enumerate() {
            match answer {
                DnsRecord::OPT { name, .. } => {
                    max_owner_len = max(max_owner_len, name.len());
                    opt = Some(answer);
                    // the iterator runs over self.answers, self.authoritative_answers and self.additional_answers,
                    // but we want the index of answer with respect to self.additional_answers
                    opt_index = i - self.answers.len() - self.authoritative_answers.len();
                },
                DnsRecord::NONOPT { name, .. } => {
                    max_owner_len = max(max_owner_len, name.len());
                }
            }
        }

        if opt.is_some() {
            additional_answers.remove(opt_index);
        }

        // Header
        res.push_str(format!("Header:\n\t{}\n\n", self.header.info_str()).as_str());

        res.push_str("Question Section:\n");
        for question in &self.questions {
            res.push('\t');
            res.push_str(question.as_padded_string(max_owner_len).as_str());
            res.push('\n');
        }
        res.push('\n');

        if !self.answers.is_empty() {
            res.push_str("Answer Section:\n");
            for answer in &self.answers {
                res.push('\t');
                res.push_str(answer.as_padded_string(max_owner_len).as_str());
                res.push('\n');
            }
            res.push('\n');
        }

        if !self.authoritative_answers.is_empty() {
            res.push_str("Authoritative Section:\n");
            for answer in &self.authoritative_answers {
                res.push('\t');
                res.push_str(answer.as_padded_string(max_owner_len).as_str());
                res.push('\n');
            }
            res.push('\n');
        }

        if !additional_answers.is_empty() {
            res.push_str("Additional Section:\n");
            for answer in &additional_answers {
                res.push('\t');
                res.push_str(answer.as_padded_string(max_owner_len).as_str());
                res.push('\n');
            }
        }

        // remove trailing '\n's
        while res.chars().nth(res.len() - 1).unwrap() == '\n' {
            res.remove(res.len() - 1);
        }
        
        write!(f, "{}", res)
    }
}
