kapiti 0.0.3

use anyhow::Result;
use bytes::BytesMut;
use lazy_static::lazy_static;
use packed_struct::prelude::*;

use crate::codec::{domain_name, message, rdata};
use crate::specs::enums_generated::{ResourceType, ResponseCode};
use crate::specs::message::{IntEnum, Message, Question, OPT};

lazy_static! {
    /// Encoder instance, currently doesn't have state
    pub static ref ENCODER: DNSMessageEncoder = DNSMessageEncoder::new();
}

/// Encodes DNS `Message`s to bytes to be sent over the wire.
pub struct DNSMessageEncoder {}

/// TTL of NXDOMAIN responses due to a domain being blocked.
/// We don't want this to be too long, in case the domain is later unblocked.
static CACHED_NXDOMAIN_TTL: u32 = 300; // 5 minutes

impl DNSMessageEncoder {
    /// Returns a `DNSMessageEncoder` for encoding DNS messages.
    pub fn new() -> DNSMessageEncoder {
        DNSMessageEncoder {}
    }

    /// Encodes a filter-based DNS response containing the provided overrides.
    /// The response could be e.g. an NXDOMAIN value for a blocked domain, or a custom A record for an overridden domain.
    pub fn encode_local_response(
        &self,
        response_code: ResponseCode,
        orig_header_id: u16,
        orig_question: &Question,
        orig_opt: &Option<OPT>,
        filter_info: &String,
        ip: Option<std::net::IpAddr>,
        udp_size_override: Option<u16>,
        buf: &mut BytesMut,
    ) -> Result<()> {
        // Cache of string offsets for use in name compression
        let mut ptr_offsets = domain_name::LabelOffsets::new();

        message::write_header_bits(
            message::HeaderBits {
                id: orig_header_id,

                is_response: true,
                op_code: Integer::from(0 /*QUERY*/),
                authoritative: false,
                truncated: false,
                recursion_desired: true,
                recursion_available: true,
                reserved_9: false,
                authentic_data: false,
                checking_disabled: false,
                response_code: Integer::from(response_code as u8),

                question_count: 1,
                answer_count: match ip {
                    Some(_) => 1, // A or AAAA resource
                    None => 0,
                },
                authority_count: match ip {
                    Some(_) => 0,
                    None => 1, // SOA resource
                },
                additional_count: match orig_opt {
                    Some(_) => 2,
                    None => 1,
                },
            },
            buf,
        )?;

        message::write_question(orig_question, buf, &mut ptr_offsets)?;

        // Write either an A/AAAA Answer resource or an SOA Authority resource, depending on ip
        match ip {
            Some(ip) => message::write_resource_fields(
                &message::ResourceFields {
                    name: orig_question.name.as_str(),
                    resource_type: orig_question.resource_type,
                    resource_class: orig_question.resource_class,
                    ttl: 300,
                    rdata: message::RDataFields::IP(ip),
                },
                buf,
                &mut ptr_offsets,
            )?,
            None => {
                message::write_resource_fields(
                    &message::ResourceFields {
                        name: orig_question.name.as_str(),
                        resource_type: IntEnum::Enum(ResourceType::SOA),
                        resource_class: orig_question.resource_class,
                        ttl: CACHED_NXDOMAIN_TTL,
                        rdata: match ip {
                            Some(ip) => message::RDataFields::IP(ip),
                            None => message::RDataFields::SOA(rdata::SOAFields {
                                mname: orig_question.name.as_str(),
                                rname: "kapiti.nickbp.com.",
                                serial: 42069,
                                refresh: CACHED_NXDOMAIN_TTL,
                                retry: 60 * 60,
                                expire: 3 * 24 * 60 * 60,
                                minimum: CACHED_NXDOMAIN_TTL,
                            }),
                        },
                    },
                    buf,
                    &mut ptr_offsets,
                )?;
            }
        }

        // Ensure OPT resource is written within the Additional section
        if let Some(opt) = orig_opt {
            message::write_opt(&opt, udp_size_override, buf)?;
        }

        message::write_resource_fields(
            &message::ResourceFields {
                name: orig_question.name.as_str(),
                resource_type: IntEnum::Enum(ResourceType::TXT),
                resource_class: orig_question.resource_class,
                ttl: CACHED_NXDOMAIN_TTL,
                rdata: message::RDataFields::TXT(filter_info),
            },
            buf,
            &mut ptr_offsets,
        )?;

        Ok(())
    }

    /// Encodes a response from the provided object representation
    pub fn encode(
        &self,
        message: &Message,
        udp_size_override: Option<u16>,
        buf: &mut BytesMut,
    ) -> Result<()> {
        // Cache of string offsets for use in name compression
        let mut ptr_offsets = domain_name::LabelOffsets::new();

        message::write_header_id(message, message.header.id, buf)?;

        for q in &message.question {
            message::write_question(&q, buf, &mut ptr_offsets)?;
        }

        for a in &message.answer {
            message::write_resource(&a, buf, &mut ptr_offsets)?;
        }

        for a in &message.authority {
            message::write_resource(a, buf, &mut ptr_offsets)?;
        }

        // Ensure OPT resource is written within the Additional section
        if let Some(opt) = &message.opt {
            message::write_opt(opt, udp_size_override, buf)?;
        }

        for a in &message.additional {
            message::write_resource(a, buf, &mut ptr_offsets)?;
        }

        Ok(())
    }
}