use super::{get_acceptable_key, remove_acceptable_key}; use crate::domain::{Action, DnatMapping, DnatMappingsPort, ModuleArgs, Record, Singleton, Value}; use crate::singleton_borrow; use chrono::{Duration, Utc}; type FieldAndValue = (Option, Option); struct DnatMappingSpec { pub src_addr: FieldAndValue, pub src_port: FieldAndValue, pub internal_addr: FieldAndValue, pub internal_port: FieldAndValue, pub dest_addr: FieldAndValue, pub dest_port: FieldAndValue, pub keep_duration: Duration, } pub struct DnatCapture { mappings: Singleton, specs: DnatMappingSpec, } impl DnatCapture { pub fn from_args(mut args: ModuleArgs, mappings: Singleton) -> DnatCapture { let src_addr = ( Some( remove_acceptable_key(&mut args, "saddr") .expect("The DnatCapture action needs a log field for the source address in “saddr”"), ), None, ); let src_port = (remove_acceptable_key(&mut args, "sport"), None); let internal_addr = ( remove_acceptable_key(&mut args, "addr"), remove_acceptable_key(&mut args, "addrValue"), ); if let &(None, None) = &internal_addr { panic!("The DnatCapture action requires either a field (“addr”) or a value (“addrValue”) for the internal address"); } let internal_port = ( remove_acceptable_key(&mut args, "port"), remove_acceptable_key(&mut args, "portValue"), ); let dest_addr = ( remove_acceptable_key(&mut args, "daddr"), remove_acceptable_key(&mut args, "daddrValue"), ); let dest_port = ( remove_acceptable_key(&mut args, "dport"), remove_acceptable_key(&mut args, "dportValue"), ); let keep_duration = match args.remove("keepSeconds") { Some(Value::Int(i)) => Duration::seconds(i as i64), _ => Duration::seconds(63), }; DnatCapture { mappings, specs: DnatMappingSpec { src_addr, src_port, internal_addr, internal_port, dest_addr, dest_port, keep_duration, }, } } } impl Action for DnatCapture { fn act(&mut self, record: &mut Record) -> Result<(), ()> { let src_addr = value_for(&self.specs.src_addr, record); let internal_addr = value_for(&self.specs.internal_addr, record); if src_addr == None || internal_addr == None { return Ok(()); } let src_port = value_for(&self.specs.src_port, record); let internal_port = value_for(&self.specs.internal_port, record); let dest_addr = value_for(&self.specs.dest_addr, record); let dest_port = value_for(&self.specs.dest_port, record); singleton_borrow!(self.mappings).put(DnatMapping { src_addr, src_port, internal_addr, internal_port, dest_addr, dest_port, keep_until: Utc::now() + self.specs.keep_duration, }); Ok(()) } } fn value_for(spec: &FieldAndValue, record: &Record) -> Option { (&spec.0) .as_deref() .and_then(|s| get_acceptable_key(record, s)) .or(spec.1.clone()) } #[cfg(test)] mod tests { use super::DnatCapture; use crate::domain::test_util::FakeDnatMappings; use crate::domain::{Action, DnatMapping, DnatMappingsPort, ModuleArgs, Record, Value}; use crate::{singleton_borrow, singleton_new, singleton_share}; use chrono::{Duration, Utc}; use std::collections::HashMap; #[test] #[should_panic( expected = "The DnatCapture action needs a log field for the source address in “saddr”" )] fn when_no_saddr_then_error() { let mut args = HashMap::with_capacity(1); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("addr".to_string(), Value::Str("int_ip".to_string())); let _ = DnatCapture::from_args(args, singleton_share!(mappings)); } #[test] #[should_panic( expected = "The DnatCapture action requires either a field (“addr”) or a value (“addrValue”) for the internal address" )] fn when_no_addr_nor_addr_value_then_error() { let mut args = HashMap::with_capacity(1); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("saddr".to_string(), Value::Str("src_ip".to_string())); let _ = DnatCapture::from_args(args, singleton_share!(mappings)); } #[test] fn when_no_addr_but_addr_value_then_no_error() { let mut args = HashMap::with_capacity(2); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("saddr".to_string(), Value::Str("src_ip".to_string())); args.insert("addrValue".to_string(), Value::Str("1.2.3.4".to_string())); let _ = DnatCapture::from_args(args, singleton_share!(mappings)); assert!(true); } #[test] fn when_no_addr_value_but_addr_then_no_error() { let mut args = HashMap::with_capacity(2); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("saddr".to_string(), Value::Str("src_ip".to_string())); args.insert("addr".to_string(), Value::Str("int_ip".to_string())); let _ = DnatCapture::from_args(args, singleton_share!(mappings)); assert!(true); } #[test] fn when_no_keep_seconds_then_63sec() { let mut args = HashMap::with_capacity(2); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("saddr".to_string(), Value::Str("src_ip".to_string())); args.insert("addr".to_string(), Value::Str("int_ip".to_string())); let action = DnatCapture::from_args(args, singleton_share!(mappings)); assert_eq!(Duration::seconds(63), action.specs.keep_duration); } #[test] fn when_insufficient_entry_then_no_mapping() { let mut args = HashMap::with_capacity(2); let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); args.insert("saddr".to_string(), Value::Str("src_ip".to_string())); args.insert("addr".to_string(), Value::Str("int_ip".to_string())); let mut action = DnatCapture::from_args(args, singleton_share!(mappings)); action.act(&mut HashMap::new()).unwrap(); assert_eq!(0, singleton_borrow!(mappings).mappings.len()); } fn when_field_and_or_value_then_check_mapping( mut args: ModuleArgs, entry_with_addr: bool, entry_with_daddr: bool, expect: DnatMapping, ) { let mappings = singleton_new!(FakeDnatMappings { mappings: Vec::new() }); // specify the Action args.insert("saddr".to_string(), Value::Str("sa".to_string())); // prepare the entry let mut entry: Record = HashMap::with_capacity(6); entry.insert("sa".to_string(), Value::Str("vsa".to_string())); entry.insert("sp".to_string(), Value::Str("vsp".to_string())); if entry_with_addr { entry.insert("a".to_string(), Value::Str("va".to_string())); entry.insert("p".to_string(), Value::Str("vp".to_string())); } if entry_with_daddr { entry.insert("da".to_string(), Value::Str("vda".to_string())); entry.insert("dp".to_string(), Value::Str("vdp".to_string())); } // run let mut action = DnatCapture::from_args(args, singleton_share!(mappings)); action.act(&mut entry).unwrap(); // check the result assert_eq!(1, singleton_borrow!(mappings).get_all().len()); let got = singleton_borrow!(mappings) .get_all() .last() .map(|m| { let mut m = (**m).clone(); m.keep_until = expect.keep_until; m }) .unwrap(); assert_eq!(expect, got); } #[test] fn when_sufficient_record_a_mapping_is_stored() { when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a")]), true, true, test_dnat_mapping(None, Some("va"), None, None, None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addrValue", "x")]), true, true, test_dnat_mapping(None, Some("x"), None, None, None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("addrValue", "x")]), true, true, test_dnat_mapping(None, Some("va"), None, None, None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("addrValue", "x")]), false, true, test_dnat_mapping(None, Some("x"), None, None, None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("daddr", "da")]), true, true, test_dnat_mapping(None, Some("va"), None, Some("vda"), None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("daddrValue", "x")]), true, true, test_dnat_mapping(None, Some("va"), None, Some("x"), None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("daddr", "da"), ("daddrValue", "x")]), true, true, test_dnat_mapping(None, Some("va"), None, Some("vda"), None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("daddr", "da"), ("daddrValue", "x")]), true, false, test_dnat_mapping(None, Some("va"), None, Some("x"), None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("port", "p")]), true, true, test_dnat_mapping(None, Some("va"), Some("vp"), None, None), ); when_field_and_or_value_then_check_mapping( as_args(vec![("addr", "a"), ("dport", "dp")]), true, true, test_dnat_mapping(None, Some("va"), None, None, Some("vdp")), ); } fn test_dnat_mapping( src_port: Option<&str>, internal_addr: Option<&str>, internal_port: Option<&str>, dest_addr: Option<&str>, dest_port: Option<&str>, ) -> DnatMapping { DnatMapping { src_addr: Some("vsa".to_string()), src_port: src_port.map(|s| s.to_string()), internal_addr: internal_addr.map(|s| s.to_string()), internal_port: internal_port.map(|s| s.to_string()), dest_addr: dest_addr.map(|s| s.to_string()), dest_port: dest_port.map(|s| s.to_string()), keep_until: Utc::now(), } } fn as_args(data: Vec<(&str, &str)>) -> ModuleArgs { let mut args = HashMap::with_capacity(data.len()); data.iter().for_each(|(f, v)| { args.insert(f.to_string(), Value::Str(v.to_string())); }); args } }