Error handling, overall code review and moved mac vendors database to a

static instance
This commit is contained in:
rzuasti
2026-01-20 14:28:31 -05:00
parent 9d5e3eaaf1
commit f9c625c7d7
7 changed files with 129 additions and 79 deletions
-6
View File
@@ -1,6 +0,0 @@
- [ ] Proper error handling across the board
- [ ] Insert new devices in database or update existing ones last_seen
- [ ] Generate an event when a new device is inserted
- [ ] Run as an endless loop with a configurable pause
- [ ] Deploy as a service in a docker container
- [ ] Deploy as a service via a nix flake
+22 -11
View File
@@ -1,6 +1,6 @@
use include_dir::{Dir, include_dir};
use log::debug;
use rusqlite::{Connection, params};
use log::{debug, error};
use rusqlite::Connection;
use rusqlite_migration::Migrations;
use std::result::Result;
use std::sync::LazyLock;
@@ -11,15 +11,26 @@ static MIGRATIONS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/database_migratio
static MIGRATIONS: LazyLock<Migrations<'static>> =
LazyLock::new(|| Migrations::from_directory(&MIGRATIONS_DIR).unwrap());
pub fn init_db() -> Result<Connection, &'static str> {
debug!("Opening database");
let mut conn = Connection::open("./oott.db").unwrap();
pub fn init_db() -> Result<Connection, String> {
debug!("Opening database.");
let mut conn = match Connection::open("./oott.db") {
Ok(value) => value,
Err(error) => {
error!("Error opening database (oott.db): {error}");
return Err(format!("Error opening database (oott.db): {error}"));
}
};
debug!("Database open, executing migrations if needed");
debug!("Database open, executing migrations if needed.");
// Update the database schema, atomically
MIGRATIONS.to_latest(&mut conn);
debug!("Database up to date");
Ok(conn)
match MIGRATIONS.to_latest(&mut conn) {
Ok(_) => {
debug!("Database up to date.");
Ok(conn)
}
Err(error) => {
error!("Error updating database: {error}");
Err(format!("Error updating database: {error}"))
}
}
}
+9 -4
View File
@@ -5,16 +5,21 @@ use std::fmt;
pub struct Device {
pub mac_address: String,
pub ipv4_address: String,
pub vendor: String,
}
impl Device {
pub fn get_mac_prefix(&self) -> String {
self.mac_address.get(0..8).unwrap_or("").to_string()
}
// pub fn get_mac_prefix(&self) -> String {
// self.mac_address.get(0..8).unwrap_or("").to_string()
// }
}
impl fmt::Display for Device {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "mac={}, ip={}", self.mac_address, self.ipv4_address)
write!(
f,
"mac={}, ip={}, vendor={}",
self.mac_address, self.ipv4_address, self.vendor
)
}
}
+42 -13
View File
@@ -1,7 +1,7 @@
mod packet_send_receive;
use crate::{config, device_finders::Device};
use log::{debug, info, warn};
use log::{debug, error, info, warn};
use packet_send_receive::{listen_for_packets, send_packet};
use pnet::{
datalink::{self, Channel, NetworkInterface},
@@ -12,20 +12,25 @@ use tokio::time::{Duration, timeout};
const DEFAULT_SENDER_TIMEOUT: u64 = 60; // 1 minute to send all packets - good for a class C network
const DEFAULT_SCAN_DURATION: u64 = 300; // 5 minutes to wait per round to receive responses
pub async fn find(interface: &str) -> Vec<Device> {
pub async fn find(interface: &str) -> Result<Vec<Device>, String> {
debug!("Looking up devices via ARP using interface {}", interface);
// Get the network device to use
let network_interface: NetworkInterface = datalink::interfaces()
let network_interface: NetworkInterface = match datalink::interfaces()
.iter()
.filter(|el| el.is_up())
.filter(|el| el.name == interface)
.next()
.expect("Selected interface not found")
.clone();
{
Some(value) => value.clone(),
None => {
error!("Interface ({interface}) not found or not active.");
return Err(format!("Interface ({interface}) not found or not active.").to_string());
}
};
// Get the IPV4 network to use
let ipv4_net = network_interface
let ipv4_net = match network_interface
.ips
.iter()
.filter_map(|el| match el {
@@ -33,13 +38,25 @@ pub async fn find(interface: &str) -> Vec<Device> {
_ => None,
})
.next()
.expect("No local ip address found");
{
Some(value) => value,
None => {
error!("No IP address found for selected interface ({interface}).");
return Err(
format!("No IP address found for selected interface ({interface}).").to_string(),
);
}
};
// Get the local MAC address
let mac = match network_interface.mac {
Some(mac) => mac,
None => {
panic!("No local MAC address found");
error!("Could not get MAC address for selected interface ({interface}).");
return Err(
format!("Could not get MAC address for selected interface ({interface}).")
.to_string(),
);
}
};
@@ -48,11 +65,20 @@ pub async fn find(interface: &str) -> Vec<Device> {
// Data channel
debug!("Creating data channel");
let tunnel = datalink::channel(&network_interface, datalink::Config::default())
.expect("Failed to create datalink channel");
let tunnel: Channel = match datalink::channel(&network_interface, datalink::Config::default()) {
Ok(value) => value,
Err(error) => {
error!("Could not create data channel: {error}");
return Err(format!("Could not create data channel: {error}").to_string());
}
};
let (sender, receiver) = match tunnel {
Channel::Ethernet(tx, rx) => (tx, rx),
_ => panic!("Unsupported data channel type"),
_ => {
error!("Unsupported data channel type");
return Err("Unsupported data channel type".to_string());
}
};
let send_interface = network_interface.clone();
@@ -68,7 +94,10 @@ pub async fn find(interface: &str) -> Vec<Device> {
info!("Receiver timeout set to {} seconds", receiver_timeout);
if sender_timeout >= scan_duration {
panic!("OOTT_ARP_SENDER_TIMEOUT needs to be smaller than OOTT_ARP_SCAN_DURATION");
warn!(
"OOTT_ARP_SENDER_TIMEOUT ({sender_timeout}) needs to be smaller than OOTT_ARP_SCAN_DURATION ({scan_duration})."
);
return Err(format!("OOTT_ARP_SENDER_TIMEOUT ({sender_timeout}) needs to be smaller than OOTT_ARP_SCAN_DURATION ({scan_duration}).").to_string());
}
let result = tokio::join!(
@@ -88,5 +117,5 @@ pub async fn find(interface: &str) -> Vec<Device> {
(_, Err(_)) => info!("ARP receiver timed out"),
};
result.1.unwrap_or(Vec::new())
Ok(result.1.unwrap_or(Vec::new()))
}
+12 -5
View File
@@ -1,4 +1,5 @@
use crate::device_finders::Device;
use crate::mac_vendor_finder;
use log::{debug, info};
use pnet::datalink::{DataLinkReceiver, DataLinkSender, NetworkInterface};
use pnet::ipnetwork::Ipv4Network;
@@ -84,14 +85,20 @@ pub async fn listen_for_packets(
if arp_packet.get_operation() == ArpOperations::Reply {
debug!("It is an ARP reply packet");
if arp_packet.get_target_proto_addr() == ipv4_net.ip() {
let packet_mac_address = arp_packet.get_sender_hw_addr().to_string();
let packet_ip_address = arp_packet.get_sender_proto_addr().to_string();
let packet_vendor = mac_vendor_finder::find(
packet_mac_address.get(0..8).unwrap_or("").to_string(),
);
debug!(
"Found online device - IP addr={} - MAC addr={}",
arp_packet.get_sender_proto_addr(),
arp_packet.get_sender_hw_addr()
"Found online device - IP addr={} - MAC addr={} - vendor={}",
packet_ip_address, packet_mac_address, packet_vendor
);
devices.push(Device {
mac_address: arp_packet.get_sender_hw_addr().to_string(),
ipv4_address: arp_packet.get_sender_proto_addr().to_string(),
mac_address: packet_mac_address,
ipv4_address: packet_ip_address,
vendor: packet_vendor,
});
}
}
+33 -24
View File
@@ -1,11 +1,8 @@
use log::{debug, info};
use log::{debug, error, info};
use serde::Deserialize;
use std::collections::HashMap;
use std::fs;
pub struct MacVendorFinder {
mac_vendors_database: HashMap<String, String>,
}
use std::sync::LazyLock;
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -17,29 +14,41 @@ struct MacRecord {
// last_update: String,
}
impl MacVendorFinder {
pub fn new() -> MacVendorFinder {
info!("Loading MAC vendor database into memory");
// Initialize mac vendors database as a static lazy loaded unit
static MAC_VENDORS_DATABASE: LazyLock<HashMap<String, String>> = LazyLock::new(|| {
info!("Loading MAC vendor database into memory");
let mut database = HashMap::new();
let data = fs::read_to_string("data/mac-vendors-export.json").unwrap();
let json: Vec<MacRecord> = serde_json::from_str(&data).unwrap();
debug!("Found {} records in the database", json.len());
for el in json {
database.insert(el.mac_prefix.to_uppercase(), el.vendor_name);
let mut database = HashMap::new();
let data = match fs::read_to_string("data/mac-vendors-export.json") {
Ok(value) => value,
Err(error) => {
error!("Error reading mac vendors database (data/mac-vendors-export.json): {error}");
panic!("Error reading mac vendors database (data/mac-vendors-export.json): {error}");
}
};
info!("MAC vendor database loaded");
MacVendorFinder {
mac_vendors_database: database,
let json: Vec<MacRecord> = match serde_json::from_str(&data) {
Ok(value) => value,
Err(error) => {
error!("Error parsing mac vendors database (data/mac-vendors-export.json): {error}");
panic!("Error parsing mac vendors database (data/mac-vendors-export.json): {error}");
}
};
debug!("Found {} records in the database", json.len());
for el in json {
database.insert(el.mac_prefix.to_uppercase(), el.vendor_name);
}
pub fn find(&mut self, mac_prefix: &str) -> Option<&String> {
self.mac_vendors_database
.get(mac_prefix.to_uppercase().as_str())
}
info!("MAC vendor database loaded");
database
});
// Find a vendor based on the MAC prefix
pub fn find(mac_prefix: String) -> String {
MAC_VENDORS_DATABASE
.get(&mac_prefix.to_uppercase())
.unwrap_or(&"".to_string())
.to_string()
}
+11 -16
View File
@@ -1,6 +1,4 @@
use crate::mac_vendor_finder::MacVendorFinder;
use log::info;
use rusqlite::Connection;
use log::{debug, info};
mod config;
mod db;
@@ -12,21 +10,18 @@ async fn main() {
env_logger::init();
info!("Starting up oott");
let mut mac_vendor_finder = MacVendorFinder::new();
let mut db_conn = db::init_db().unwrap();
let db_conn = db::init_db();
// let devices = device_finders::arp::find("eno1").await;
let devices = device_finders::arp::find("eno1").await.unwrap();
// info!("Done with ARP probes");
// info!("Found {} online devices", devices.iter().count());
// for device in devices.iter() {
// info!("Device found {}", device);
// let vendor_result = mac_vendor_finder.find(device.get_mac_prefix().as_str());
// match vendor_result {
// Some(vendor) => info!("Device vendor = {}", vendor),
// None => info!("Device vendor not found"),
// }
// }
info!("Done with ARP probes");
info!("Found {} online devices", devices.iter().count());
// mac_vendor_finder.populate_vendors(&devices);
for device in devices.iter() {
debug!("Device found {}", device);
}
info!("Exiting");
}