~kline/firebee

c9eb0e2b26c9e4ca5324e0daa9d92bbcca5f53d5 — Gareth Pulham 3 months ago 9565e75
Basic DNS server
2 files changed, 105 insertions(+), 2 deletions(-)

M Cargo.toml
M src/main.rs
M Cargo.toml => Cargo.toml +2 -0
@@ 6,3 6,5 @@ edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
domain = "0.6.1"
itertools = "0.10.1"

M src/main.rs => src/main.rs +103 -2
@@ 1,3 1,104 @@
fn main() {
    println!("Hello, world!");
use std::net::{UdpSocket, IpAddr, AddrParseError};
use std::error::Error;
use itertools::Itertools;
use domain::base::{Message, MessageBuilder, Dname, ToDname};
use domain::base::name::RelativeDname;
use domain::base::iana::{Class, Rcode};
use domain::rdata::A;

fn is_banned(ip: IpAddr) -> Option<u8> {
    let octet: u8 = match ip {
        IpAddr::V4(address) => {
            *address.octets().last().unwrap()
        },
        IpAddr::V6(address) => {
            *address.octets().last().unwrap()
        }
    };
    Some(octet)
}

fn parse_ip(name: RelativeDname<Vec<u8>>) -> Result<IpAddr, AddrParseError> {
    match name.label_count() {
        4 => {
            name.iter()
                .rev()
                .map(|l| l.to_string())
                .collect::<Vec<String>>()
                .join(".")
                .parse::<IpAddr>()
        },
        32 => {
            let addr_nibbles = name.iter()
                .rev()
                .map(|l| l.to_string())
                .collect::<Vec<String>>();

            addr_nibbles.chunks(4)
                .into_iter()
                .map(|pair| pair.to_vec().join("") )
                .join(":")
                .parse::<IpAddr>()
        },
        _ => {
            // There's no way to manually raise an AddrParseError here
            // since the stdlib arbitrarily marks it private for no reason.
            // Here's how we do it instead.
            "IP address mis-sized".parse::<IpAddr>()
        }
    }
}

fn main() -> std::io::Result<()> {
    let socket = UdpSocket::bind("[::]:5252")?;

    loop {
        // Get the next UDP packet
        let mut buf = vec![0; 512];
        let (_amt, src) = socket.recv_from(&mut buf)?;

        // Try to parse it as a DNS packet
        let packet: Message<Vec<u8>> =  match Message::from_octets(buf.into()) {
            Ok(packet) => packet,
            Err(_) => continue,
        };

        let response = match dns_banned(&packet) {
            Ok(res) => {
                match res {
                    Some(ban_code) => {
                        let mut response = MessageBuilder::new_vec()
                            .start_answer(&packet, Rcode::NoError)
                            .unwrap();
                        response.push(
                            (packet.sole_question().unwrap().qname(), Class::In, 10, A::from_octets(127,0,0,ban_code))
                        ).unwrap();
                        response
                    },
                    None => MessageBuilder::new_vec().start_answer(&packet, Rcode::NXDomain).unwrap(),
                }
            },
            Err(_) => {
                MessageBuilder::new_vec()
                    .start_answer(&packet, Rcode::Refused)
                    .unwrap()
            }
        };
        let _ = socket.send_to(&response.finish(), &src);
    }
}
fn dns_banned(packet: &Message<Vec<u8>>) -> Result<Option<u8>, Box<dyn Error>> {
    // Get the first question
    let question = packet.sole_question()?;

    // Convert that question to an IpAddr
    let hostname: Dname<Vec<u8>> = ToDname::to_dname(question.qname())?;
    let address = match hostname.strip_suffix(&Dname::vec_from_str("kernel.panic.")?) {
        Ok(address) => address,
        Err(_) => { Err("Couldn't strip suffix")? }
    };
    let ip = parse_ip(address)?;

    // Check if they are banned
    Ok(is_banned(ip))
}