kapiti 0.0.3

use crate::specs::{enums_generated, rdata};

use bytecheck::CheckBytes;
use rkyv::{Archive, Deserialize, Serialize};
use std::fmt;
use std::net::{Ipv4Addr, Ipv6Addr};

// Object renderings of DNS requests and responses as they would be sent/received over the wire.
// Some minimal preprocessing is performed vs the raw DNS data:
// - Resource count and string length fields are omitted here since we have them via the objects.
// - Compressed names are decompressed/dereferenced as any byte pointers wouldn't be relevant anyway.

/// A DNS message, either a request or a response.
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct Message {
    pub header: Header,

    /// Moved to the root Message to simplify handling.
    /// The wire format meanwhile serializes this as an Additional Resource.
    pub opt: Option<OPT>,

    pub question: Vec<Question>,
    pub answer: Vec<Resource>,
    pub authority: Vec<Resource>,
    pub additional: Vec<Resource>,
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Message:")?;
        write!(f, "\n- header: {}", self.header)?;
        if let Some(opt) = &self.opt {
            write!(f, "\n- opt: {}", opt)?;
        }
        for q in &self.question {
            write!(f, "\n- question: {}", q)?;
        }
        for a in &self.answer {
            write!(f, "\n- answer: {}", a)?;
        }
        for a in &self.authority {
            write!(f, "\n- authority: {}", a)?;
        }
        for a in &self.additional {
            write!(f, "\n- additional: {}", a)?;
        }
        Ok(())
    }
}

/// Wrapper for an enum that can fall back to being an int.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, Copy, PartialEq)]
#[archive_attr(derive(CheckBytes))]
pub enum IntEnum<T, E> {
    Enum(E),
    Unknown(T),
}

// From RFC2535 section 6.1 or RFC2929 section 2
// (two bits added following RFC1035)
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      ID                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |QR|   Opcode  |AA|TC|RD|RA| Z|AD|CD|   RCODE   |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    QDCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ANCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    NSCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ARCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct Header {
    /// Random request ID, 16 bits
    pub id: u16,

    /// Whether this is a request (false) or a response (true)
    pub is_response: bool,

    /// Raw Opcode value, 4 bits
    pub op_code: IntEnum<u8, enums_generated::OpCode>, // 4 bits

    /// ref: https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-12

    /// RFC1035: AA/Authoritative Answer (response-only)
    pub authoritative: bool,

    /// RFC1035: TC/Truncated Response (response-only)
    pub truncated: bool,

    /// RFC1035: RD/Recursion Desired
    pub recursion_desired: bool,

    /// RFC1035: RA/Recursion Available (response-only)
    pub recursion_available: bool,

    /// bit 9: Reserved
    pub reserved_9: bool,

    /// RFC4035/RFC6840/RFC Errata 4924: AD/Authentic Data
    pub authentic_data: bool,

    /// RFC4035/RFC6840/RFC Errata 4927: CD/Checking Disabled
    pub checking_disabled: bool,

    /// Raw response_code value, 4 bits
    pub response_code: IntEnum<u8, enums_generated::ResponseCode>, // 4 bits
}

impl fmt::Display for Header {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "id={} is_response={} op_code=",
            self.id, self.is_response
        )?;
        match &self.op_code {
            IntEnum::Enum(e) => write!(f, "{:?}", e)?,
            IntEnum::Unknown(i) => write!(f, "{}", i)?,
        }
        write!(
            f,
            " authoritative={} truncated={} recursion_desired={} recursion_available={} authentic_data={} checking_disabled={} response_code=",
            self.authoritative, self.truncated, self.recursion_desired, self.recursion_available, self.authentic_data, self.checking_disabled,
        )?;
        match &self.response_code {
            IntEnum::Enum(e) => write!(f, "{:?}", e)?,
            IntEnum::Unknown(i) => write!(f, "{}", i)?,
        }
        Ok(())
    }
}

/// OPT rdata, serialized as an Additional Resource but represented here as root-level header data to simplify processing.
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct OPT {
    /// EDNS(0) header values from [RFC6891]:

    /// 16 bit value from CLASS field
    /// The requestor's UDP payload size (encoded in the RR CLASS field) is
    /// the number of octets of the largest UDP payload that can be
    /// reassembled and delivered in the requestor's network stack.
    pub udp_size: u16,

    /// first 8 bits from TTL field
    /// Forms the upper 8 bits of extended 12-bit RCODE, together with the
    /// 4 bits defined in [RFC1035].
    pub response_code: u8,

    /// second 8 bits from TTL field
    /// Indicates the implementation level of the setter, currently '0'
    pub version: u8,

    /// bit 16 from TTL field
    /// DNSSEC OK bit as defined by [RFC3225].
    pub dnssec_ok: bool,

    pub option: Vec<OPTOption>,
}

impl fmt::Display for OPT {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "udp_size={} response_code={}",
            self.udp_size, self.response_code
        )?;
        write!(
            f,
            " version={} dnssec_ok={} option=[",
            self.version, self.dnssec_ok
        )?;
        for entry in &self.option {
            write!(f, " (code=")?;
            match &entry.code {
                IntEnum::Enum(e) => write!(f, "{:?}", e)?,
                IntEnum::Unknown(i) => write!(f, "{}", i)?,
            }
            write!(f, " data={:02X?})", entry.data)?;
        }
        write!(f, " ]")?;
        Ok(())
    }
}

/// An Option entry from OPT rdata.
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct OPTOption {
    /// Assigned by the Expert Review process as defined by the DNSEXT
    /// working group and the IESG.
    pub code: IntEnum<u16, enums_generated::OPTOptionCode>,

    pub data: Vec<u8>,
}

// From RFC1035 section 4.1.2
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                     QNAME                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QTYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QCLASS                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct Question {
    /// DNS name that this question is referring to, variable length
    pub name: String,

    /// Resource type value, 16 bits
    pub resource_type: IntEnum<u16, enums_generated::ResourceType>,

    /// Resource class value, 16 bits
    pub resource_class: IntEnum<u16, enums_generated::ResourceClass>,
}

impl fmt::Display for Question {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "name={} resource_type=", self.name)?;
        match &self.resource_type {
            IntEnum::Enum(e) => write!(f, "{:?}", e)?,
            IntEnum::Unknown(i) => write!(f, "{}", i)?,
        }
        write!(f, " resource_class=")?;
        match &self.resource_class {
            IntEnum::Enum(e) => write!(f, "{:?}", e)?,
            IntEnum::Unknown(i) => write!(f, "{}", i)?,
        }
        Ok(())
    }
}

// From RFC1035 section 4.1.3
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                                               /
// /                      NAME                     /
// |                                               |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     CLASS                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TTL                      |
// |                                               |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                   RDLENGTH                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--|
// /                     RDATA                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub struct Resource {
    /// DNS name that this resource is referring to, variable length
    pub name: String,

    /// Resource type value, 16 bits
    pub resource_type: IntEnum<u16, enums_generated::ResourceType>,

    /// Resource class value, 16 bits.
    pub resource_class: IntEnum<u16, enums_generated::ResourceClass>,

    /// Time to live in seconds, 32 bits.
    pub ttl: u32,

    /// Resource-specific data, variable length
    pub rdata: ResourceData,
}

impl fmt::Display for Resource {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "name={} resource_type=", self.name)?;
        match &self.resource_type {
            IntEnum::Enum(e) => write!(f, "{:?}", e)?,
            IntEnum::Unknown(i) => write!(f, "{}", i)?,
        }
        // Ignore inapplicable resource_class/ttl for OPT records
        if self.resource_type != IntEnum::Enum(enums_generated::ResourceType::OPT) {
            write!(f, " resource_class=")?;
            match &self.resource_class {
                IntEnum::Enum(e) => write!(f, "{:?}", e)?,
                IntEnum::Unknown(i) => write!(f, "{}", i)?,
            }
            write!(f, " ttl={}", self.ttl)?;
        }
        write!(f, " {}", self.rdata)
    }
}

/// Parsed (or raw) content from the Resource RDATA field.
/// If the Resource type is supported, the data is parsed into a message.
/// Otherwise it is provided in its original byte form via raw_data.
#[derive(Archive, Deserialize, Serialize, Clone)]
#[archive_attr(derive(CheckBytes))]
pub enum ResourceData {
    /// If the resource type is unsupported, it's passed through in raw form.
    RawData(Vec<u8>),

    // Supported resource types, covering items from: https://en.wikipedia.org/wiki/List_of_DNS_record_types
    // For descriptions of each type, see the more exhaustive list in enums_generated.rs
    // Note that OPT data is not located here, instead it's at the root Message.
    A(rdata::A),                   // RFC1035
    NS(rdata::NS),                 // RFC1035
    CNAME(rdata::CNAME),           // RFC1035
    SOA(rdata::SOA),               // RFC1035, RFC2308
    PTR(rdata::PTR),               // RFC1035
    HINFO(rdata::HINFO),           // RFC1035
    MX(rdata::MX),                 // RFC1035, RFC7505
    TXT(rdata::TXT),               // RFC1035
    RP(rdata::RP),                 // RFC1183
    AFSDB(rdata::AFSDB),           // RFC1183
    SIG(rdata::SIG),               // RFC2535
    KEY(rdata::KEY),               // RFC2535, RFC2930
    AAAA(rdata::AAAA),             // RFC3596
    LOC(rdata::LOC),               // RFC1876
    SRV(rdata::SRV),               // RFC2782
    NAPTR(rdata::NAPTR),           // RFC3403
    KX(rdata::KX),                 // RFC2230
    CERT(rdata::CERT),             // RFC4398
    DNAME(rdata::DNAME),           // RFC6672
    APL(rdata::APL),               // RFC3123
    DS(rdata::DS),                 // RFC4034. also:  CDS/59, TS/32768, DLV/32769
    SSHFP(rdata::SSHFP),           // RFC4255
    IPSECKEY(rdata::IPSECKEY),     // RFC4025
    RRSIG(rdata::RRSIG),           // RFC4034
    NSEC(rdata::NSEC),             // RFC4034
    DNSKEY(rdata::DNSKEY),         // RFC4034, also:  CDNSKEY/60
    DHCID(rdata::DHCID),           // RFC4701
    NSEC3(rdata::NSEC3),           // RFC5155
    NSEC3PARAM(rdata::NSEC3PARAM), // RFC5155
    TLSA(rdata::TLSA),             // RFC6698
    SMIMEA(rdata::SMIMEA),         // RFC8162
    HIP(rdata::HIP),               // RFC8005
    OPENPGPKEY(rdata::OPENPGPKEY), // RFC7929
    CSYNC(rdata::CSYNC),           // RFC7477
    TKEY(rdata::TKEY),             // RFC2930
    TSIG(rdata::TSIG),             // RFC2845
    URI(rdata::URI),               // RFC7553
    CAA(rdata::CAA),               // RFC6844
}

impl fmt::Display for ResourceData {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        return match self {
            ResourceData::RawData(raw_data) => write!(f, "raw_data={:02X?}", raw_data),
            ResourceData::A(data) => write!(
                f,
                "a[address={}]",
                Ipv4Addr::new(
                    data.address1,
                    data.address2,
                    data.address3,
                    data.address4,
                )
            ),
            ResourceData::NS(data) => write!(f, "ns[nsdname={:?}]", data.nsdname),
            ResourceData::CNAME(data) => write!(f, "cname[cname={:?}]", data.cname),
            ResourceData::SOA(data) => write!(
                f,
                "soa=[mname={:?} rname={:?} serial={} refresh={} retry={} expire={} minimum={}]",
                data.mname,
                data.rname,
                data.serial,
                data.refresh,
                data.retry,
                data.expire,
                data.minimum,
            ),
            ResourceData::PTR(data) => write!(f, "ptr[ptrdname={:?}]", data.ptrdname),
            ResourceData::HINFO(data) => write!(f, "hinfo[cpu={:02X?} os={:02X?}]", data.cpu, data.os),
            ResourceData::MX(data) => write!(
                f,
                "mx[preference={} exchange={:?}]",
                data.preference,
                data.exchange,
            ),
            ResourceData::TXT(data) => {
                write!(f, "txt[")?;
                for e in &data.entries {
                    write!(f, " (data={:02X?})", e.data)?;
                }
                write!(f, " ]")
            },
            ResourceData::RP(data) => write!(
                f,
                "rp[mbox_dname={:?} txt_dname={:?}]",
                data.mbox_dname,
                data.txt_dname,
            ),
            ResourceData::AFSDB(data) => write!(
                f,
                "afsdb[subtype={} hostname={:?}]",
                data.subtype,
                data.hostname,
            ),
            ResourceData::SIG(data) => write!(
                f,
                "sig[type_covered={} algorithm={} labels={} original_ttl={} signature_expiration={} signature_inception={} key_tag={} signers_name={:?} signature={:02X?}]",
                data.type_covered, data.algorithm, data.labels, data.original_ttl, data.signature_expiration, data.signature_inception, data.key_tag, data.signers_name, data.signature,
            ),
            ResourceData::KEY(data) => write!(
                f,
                "key[flags={} protocol={} algorithm={} public_key={:02X?}]",
                data.flags,
                data.protocol,
                data.algorithm,
                data.public_key,
            ),
            ResourceData::AAAA(data) => write!(
                f,
                "aaaa[address={}]",
                Ipv6Addr::new(
                    data.address1,
                    data.address2,
                    data.address3,
                    data.address4,
                    data.address5,
                    data.address6,
                    data.address7,
                    data.address8,
                )
            ),
            ResourceData::LOC(data) => write!(
                f,
                "loc[version={} size={} horiz_pre={} vert_pre={} latitude={} longitude={} altitude={}]",
                data.version, data.size, data.horiz_pre, data.vert_pre, data.latitude, data.longitude, data.altitude,
            ),
            ResourceData::SRV(data) => write!(
                f,
                "srv[priority={} weight={} port={} target={:?}]",
                data.priority,
                data.weight,
                data.port,
                data.target,
            ),
            ResourceData::NAPTR(data) => write!(
                f,
                "naptr[order={} preference={} flags={:02X?} services={:02X?} regexp={:02X?} replacement={:?}]",
                data.order, data.preference, data.flags, data.services, data.regexp, data.replacement,
            ),
            ResourceData::KX(data) => write!(
                f,
                "kx[preference={} exchanger={:?}]",
                data.preference,
                data.exchanger,
            ),
            ResourceData::CERT(data) => write!(
                f,
                "cert[type={} key_tag={} algorithm={} certificate={:02X?}]",
                data.type_,
                data.key_tag,
                data.algorithm,
                data.certificate,
            ),
            ResourceData::DNAME(data) => write!(f, "dname[dname={:?}]", data.dname),
            ResourceData::APL(data) => {
                write!(f, "apl[")?;
                for entry in &data.item {
                    write!(
                        f,
                        " (address_family={} prefix={} negation={} afd_part={:02X?})",
                        entry.address_family,
                        entry.prefix,
                        entry.negation,
                        entry.afd_part,
                    )?;
                }
                write!(f, " ]")
            },
            ResourceData::DS(data) => write!(
                f,
                "ds[key_tag={} algorithm={} digest_type={} digest={:02X?}]",
                data.key_tag,
                data.algorithm,
                data.digest_type,
                data.digest,
            ),
            ResourceData::SSHFP(data) => write!(
                f,
                "sshfp[algorithm={} fp_type={} fingerprint={:02X?}]",
                data.algorithm,
                data.fp_type,
                data.fingerprint,
            ),
            ResourceData::IPSECKEY(data) => write!(
                f,
                "ipseckey[precedence={} gateway_type={} algorithm={} gateway={:02X?} public_key={:02X?}]",
                data.precedence, data.gateway_type, data.algorithm, data.gateway, data.public_key,
            ),
            ResourceData::RRSIG(data) => write!(
                f,
                "rrsig[type_covered={} algorithm={} labels={} original_ttl={} signature_expiration={} signature_inception={} key_tag={} signers_name={:?} signature={:02X?}]",
                data.type_covered, data.algorithm, data.labels, data.original_ttl, data.signature_expiration, data.signature_inception, data.key_tag, data.signers_name, data.signature,
            ),
            ResourceData::NSEC(data) => write!(
                f,
                "nsec[next_domain_name={:?} type_bit_maps={:02X?}]",
                data.next_domain_name,
                data.type_bit_maps,
            ),
            ResourceData::DNSKEY(data) => write!(
                f,
                "dnskey[flags={} protocol={} algorithm={} public_key={:02X?}]",
                data.flags,
                data.protocol,
                data.algorithm,
                data.public_key,
            ),
            ResourceData::DHCID(data) => write!(f, "dhcid[dhcp_data={:02X?}]", data.dhcp_data),
            ResourceData::NSEC3(data) => write!(
                f,
                "nsec3[hash_algorithm={} flags={} iterations={} salt={:02X?} next_hashed_owner_name={:02X?} type_bit_maps={:02X?}]",
                data.hash_algorithm, data.flags, data.iterations, data.salt, data.next_hashed_owner_name, data.type_bit_maps,
            ),
            ResourceData::NSEC3PARAM(data) => write!(
                f,
                "nsec3param[hash_algorithm={} flags={} iterations={} salt={:02X?}]",
                data.hash_algorithm,
                data.flags,
                data.iterations,
                data.salt,
            ),
            ResourceData::TLSA(data) => write!(
                f,
                "tlsa[certificate_usage={} selector={} matching_type={} certificate_association_data={:02X?}]",
                data.certificate_usage, data.selector, data.matching_type, data.certificate_association_data,
            ),
            ResourceData::SMIMEA(data) => write!(
                f,
                "smimea[certificate_usage={} selector={} matching_type={} certificate_association_data={:02X?}]",
                data.certificate_usage, data.selector, data.matching_type, data.certificate_association_data,
            ),
            ResourceData::HIP(data) => {
                write!(
                    f,
                    "hip[public_key_algorithm={} hit={:02X?} public_key={:02X?} rendezvous_servers=[",
                    data.public_key_algorithm,
                    data.hit,
                    data.public_key,
                )?;
                for rs in &data.rendezvous_servers {
                    write!(f, " '{:?}'", rs)?;
                }
                write!(f, " ]")
            },
            ResourceData::OPENPGPKEY(data) => write!(f, "openpgpkey[public_key={:?}]", data.public_key),
            ResourceData::CSYNC(data) => write!(
                f,
                "csync[soa_serial={} flags={} type_bit_map={:02X?}]",
                data.soa_serial,
                data.flags,
                data.type_bit_map,
            ),
            ResourceData::TKEY(data) => write!(
                f,
                "tkey[algorithm={:?} inception={} expiration={} mode={} error={} key_data={:02X?} other_data={:02X?}]",
                data.algorithm, data.inception, data.expiration, data.mode, data.error, data.key_data, data.other_data,
            ),
            ResourceData::TSIG(data) => write!(
                f,
                "tsig[algorithm_name={:?} time_signed={} fudge={} mac={:02X?} other_data={:02X?}]",
                data.algorithm_name,
                data.time_signed,
                data.fudge,
                data.mac,
                data.other_data,
            ),
            ResourceData::URI(data) => write!(
                f,
                "uri[priority={} weight={} target={:?}]",
                data.priority,
                data.weight,
                data.target,
            ),
            ResourceData::CAA(data) => write!(
                f,
                "caa[flags={} tag={:?} value={:02X?}]",
                data.flags,
                data.tag,
                data.value,
            ),
        };
    }
}