1 files changed, 130 insertions, 0 deletions
diff --git a/src/packet/mod.rs b/src/packet/mod.rs
new file mode 100644
index 0000000..0b7cb7b
--- /dev/null
+++ b/src/packet/mod.rs
@@ -0,0 +1,130 @@
+use std::net::IpAddr;
+
+use self::{
+    buffer::PacketBuffer, header::DnsHeader, query::QueryType, question::DnsQuestion,
+    record::DnsRecord,
+};
+
+type Error = Box<dyn std::error::Error>;
+pub type Result<T> = std::result::Result<T, Error>;
+
+pub mod buffer;
+pub mod header;
+pub mod query;
+pub mod question;
+pub mod record;
+pub mod result;
+
+#[derive(Clone, Debug)]
+pub struct Packet {
+    pub header: DnsHeader,
+    pub questions: Vec<DnsQuestion>,
+    pub answers: Vec<DnsRecord>,
+    pub authorities: Vec<DnsRecord>,
+    pub resources: Vec<DnsRecord>,
+}
+
+impl Packet {
+    pub fn new() -> Self {
+        Self {
+            header: DnsHeader::new(),
+            questions: Vec::new(),
+            answers: Vec::new(),
+            authorities: Vec::new(),
+            resources: Vec::new(),
+        }
+    }
+
+    pub fn from_buffer(buffer: &mut PacketBuffer) -> Result<Self> {
+        let mut result = Self::new();
+        result.header.read(buffer)?;
+
+        for _ in 0..result.header.questions {
+            let mut question = DnsQuestion::new("".to_string(), QueryType::UNKNOWN(0));
+            question.read(buffer)?;
+            result.questions.push(question);
+        }
+
+        for _ in 0..result.header.answers {
+            let rec = DnsRecord::read(buffer)?;
+            result.answers.push(rec);
+        }
+        for _ in 0..result.header.authoritative_entries {
+            let rec = DnsRecord::read(buffer)?;
+            result.authorities.push(rec);
+        }
+        for _ in 0..result.header.resource_entries {
+            let rec = DnsRecord::read(buffer)?;
+            result.resources.push(rec);
+        }
+
+        Ok(result)
+    }
+
+    pub fn write(&mut self, buffer: &mut PacketBuffer) -> Result<()> {
+        self.header.questions = self.questions.len() as u16;
+        self.header.answers = self.answers.len() as u16;
+        self.header.authoritative_entries = self.authorities.len() as u16;
+        self.header.resource_entries = self.resources.len() as u16;
+
+        self.header.write(buffer)?;
+
+        for question in &self.questions {
+            question.write(buffer)?;
+        }
+        for rec in &self.answers {
+            rec.write(buffer)?;
+        }
+        for rec in &self.authorities {
+            rec.write(buffer)?;
+        }
+        for rec in &self.resources {
+            rec.write(buffer)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn get_random_a(&self) -> Option<IpAddr> {
+        self.answers
+            .iter()
+            .filter_map(|record| match record {
+                DnsRecord::A { addr, .. } => Some(IpAddr::V4(*addr)),
+                DnsRecord::AAAA { addr, .. } => Some(IpAddr::V6(*addr)),
+                _ => None,
+            })
+            .next()
+    }
+
+    fn get_ns<'a>(&'a self, qname: &'a str) -> impl Iterator<Item = (&'a str, &'a str)> {
+        self.authorities
+            .iter()
+            .filter_map(|record| match record {
+                DnsRecord::NS { domain, host, .. } => Some((domain.as_str(), host.as_str())),
+                _ => None,
+            })
+            .filter(move |(domain, _)| qname.ends_with(*domain))
+    }
+
+    pub fn get_resolved_ns(&self, qname: &str) -> Option<IpAddr> {
+        self.get_ns(qname)
+            .flat_map(|(_, host)| {
+                self.resources
+                    .iter()
+                    .filter_map(move |record| match record {
+                        DnsRecord::A { domain, addr, .. } if domain == host => {
+                            Some(IpAddr::V4(*addr))
+                        }
+                        DnsRecord::AAAA { domain, addr, .. } if domain == host => {
+                            Some(IpAddr::V6(*addr))
+                        }
+                        _ => None,
+                    })
+            })
+            .next()
+    }
+
+    pub fn get_unresolved_ns<'a>(&'a self, qname: &'a str) -> Option<&'a str> {
+        self.get_ns(qname).map(|(_, host)| host).next()
+    }
+}