althea_kernel_interface/
lib.rs

1#[macro_use]
2extern crate derive_error;
3
4use std::str;
5extern crate hwaddr;
6extern crate regex;
7
8use std::collections::HashMap;
9use std::net::IpAddr;
10use hwaddr::HwAddr;
11use std::str::FromStr;
12use regex::Regex;
13use std::process::{Command, Output, ExitStatus};
14use std::os::unix::process::ExitStatusExt;
15
16#[derive(Debug, Error)]
17pub enum Error {
18    Io(std::io::Error),
19    UTF8(std::string::FromUtf8Error),
20    ParseInt(std::num::ParseIntError),
21    AddrParse(std::net::AddrParseError),
22    #[error(msg_embedded, no_from, non_std)] RuntimeError(String),
23}
24
25pub struct KernelInterface<'a> {
26    fake_outputs: Option<HashMap<(&'a str, &'a [&'a str]), Output>>,
27}
28
29impl<'a> KernelInterface<'a> {
30    pub fn new() -> KernelInterface<'a> {
31        KernelInterface {
32            fake_outputs: None
33        }
34    }
35
36    fn add_fake_outputs(&mut self, fake_outputs: HashMap<(&'a str, &'a [&'a str]), Output>) {
37        self.fake_outputs = Some(fake_outputs);
38    }
39
40    fn run_command(self, program: &'a str, args: &'a [&str]) -> Result<Output, Error> {
41        match self.fake_outputs {
42            Some(outputs) => Ok(outputs[&(program, args)].clone()),
43            None => Command::new(program).args(args).output().map_err(|e| Error::Io(e))
44        }
45    }
46
47    fn get_neighbors_linux(self) -> Result<Vec<(HwAddr, IpAddr)>, Error> {
48        let output = self.run_command("ip", &["neighbor"])?;
49        let mut vec = Vec::new();
50        let re = Regex::new(r"(\S*) .* (\S*) (REACHABLE|STALE|DELAY)").unwrap();
51        for caps in re.captures_iter(&String::from_utf8(output.stdout)?) {
52            vec.push((
53                caps.get(2).unwrap().as_str().parse::<HwAddr>()?,
54                IpAddr::from_str(&caps[1])?,
55            ));
56        }
57        Ok(vec)
58    }
59    
60    /// Returns a vector of neighbors reachable over layer 2, giving the hardware
61    /// and IP address of each. Implemented with `ip neighbor` on Linux.
62    pub fn get_neighbors(self) -> Result<Vec<(HwAddr, IpAddr)>, Error> {
63        if cfg!(target_os = "linux") {
64            return self.get_neighbors_linux();
65        }
66
67        Err(Error::RuntimeError(String::from("not implemented for this platform")))
68    }
69
70    fn get_traffic_linux(self) -> Result<Vec<(HwAddr, IpAddr, u64)>, Error> {
71        let output = self.run_command("ebtables", &["-L", "INPUT", "--Lc"])?;
72        let mut vec = Vec::new();
73        let re = Regex::new(r"-s (.*) --ip6-dst (.*)/.* bcnt = (.*)").unwrap();
74        for caps in re.captures_iter(&String::from_utf8(output.stdout)?) {
75            vec.push((
76                caps[1].parse::<HwAddr>()?,
77                IpAddr::from_str(&caps[2])?,
78                caps[3].parse::<u64>()?,
79            ));
80        }
81        Ok(vec)
82    }
83
84    /// Returns a vector of traffic coming from a specific hardware address and going
85    /// to a specific IP. Note that this will only track flows that have already been
86    /// registered. Implemented with `ebtables` on Linux.
87    pub fn get_traffic(self) -> Result<Vec<(HwAddr, IpAddr, u64)>, Error> {
88        if cfg!(target_os = "linux") {
89            return self.get_traffic_linux();
90        }
91
92        Err(Error::RuntimeError(String::from("not implemented for this platform")))
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    #[test]
100    fn test_get_neighbors_linux() {
101        let mut outputs = HashMap::new();
102        outputs.insert(
103            ("ip", &["neighbor"][..]),
104            Output {
105                stdout: b"10.0.2.2 dev eth0 lladdr 00:00:00:aa:00:03 STALE
10610.0.0.2 dev eth0  FAILED
10710.0.1.2 dev eth0 lladdr 00:00:00:aa:00:05 REACHABLE
1082001::2 dev eth0 lladdr 00:00:00:aa:00:56 REACHABLE
109fe80::7459:8eff:fe98:81 dev eth0 lladdr 76:59:8e:98:00:81 STALE
110fe80::433:25ff:fe8c:e1ea dev eth0 lladdr 1a:32:06:78:05:0a STALE
1112001::2 dev eth0  FAILED"
112                    .to_vec(),
113                stderr: b"".to_vec(),
114                status: ExitStatus::from_raw(0),
115            },
116        );
117  
118        let mut ki = KernelInterface::new();
119        ki.add_fake_outputs(outputs);
120
121        let addresses = ki.get_neighbors_linux().unwrap();
122
123        assert_eq!(format!("{}", addresses[0].0), "0:0:0:AA:0:3");
124        assert_eq!(format!("{}", addresses[0].1), "10.0.2.2");
125
126        assert_eq!(format!("{}", addresses[1].0), "0:0:0:AA:0:5");
127        assert_eq!(format!("{}", addresses[1].1), "10.0.1.2");
128
129        assert_eq!(format!("{}", addresses[2].0), "0:0:0:AA:0:56");
130        assert_eq!(format!("{}", addresses[2].1), "2001::2");
131    }
132
133    #[test]
134    fn test_get_traffic_linux() {
135        let mut outputs = HashMap::new();
136        outputs.insert(("ebtables", &["-L", "INPUT", "--Lc"][..]), Output {
137                    stdout: b"Bridge table: filter
138
139Bridge chain: INPUT, entries: 3, policy: ACCEPT
140-p IPv6 -s 0:0:0:aa:0:2 --ip6-dst 2001::1/ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff -j ACCEPT , pcnt = 1199 -- bcnt = 124696
141-p IPv6 -s 0:0:0:aa:0:0 --ip6-dst 2001::3/ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff -j ACCEPT , pcnt = 1187 -- bcnt = 123448
142-p IPv6 -s 0:0:0:aa:0:0 --ip6-dst 2001::3/ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff -j ACCEPT , pcnt = 0 -- bcnt = 0".to_vec(),
143                    stderr: b"".to_vec(),
144                    status: ExitStatus::from_raw(0)
145                });
146
147        let mut ki = KernelInterface::new();
148        ki.add_fake_outputs(outputs);
149
150        let traffic = ki.get_traffic_linux().unwrap();
151
152        assert_eq!(format!("{}", traffic[0].0), "0:0:0:AA:0:2");
153        assert_eq!(format!("{}", traffic[0].1), "2001::1");
154        assert_eq!(traffic[0].2, 124696);
155
156        assert_eq!(format!("{}", traffic[1].0), "0:0:0:AA:0:0");
157        assert_eq!(format!("{}", traffic[1].1), "2001::3");
158        assert_eq!(traffic[1].2, 123448);
159    }
160
161
162}