use std::hash::Hash;
use std::time::Duration;
use crate::analysis::{AnalyzedFlow, L7Summary};
use crate::correlate::KeyIndexed;
use crate::detect::patterns::DgaScorer;
use crate::detect::risk::{is_weak_cipher, label_matches};
use crate::detect::{FlowRisk, IocKind, IocMatch, IocSet};
use crate::{FlowStats, KeyFields, Timestamp};
pub struct FlowAnalyzer<K>
where
K: Hash + Eq,
{
summaries: KeyIndexed<K, L7Summary>,
ioc: Option<IocSet>,
dga: DgaScorer,
}
impl<K> FlowAnalyzer<K>
where
K: Hash + Eq + Clone,
{
pub fn new(ttl: Duration) -> Self {
Self {
summaries: KeyIndexed::new_unbounded(ttl),
ioc: None,
dga: DgaScorer::new(),
}
}
pub fn with_capacity(ttl: Duration, capacity: usize) -> Self {
Self {
summaries: KeyIndexed::new(ttl, capacity),
ioc: None,
dga: DgaScorer::new(),
}
}
pub fn with_ioc(mut self, ioc: IocSet) -> Self {
self.ioc = Some(ioc);
self
}
pub fn len(&self) -> usize {
self.summaries.len()
}
pub fn is_empty(&self) -> bool {
self.summaries.is_empty()
}
#[cfg(any(feature = "tls", feature = "http", feature = "dns"))]
fn summary_mut(&mut self, key: &K, ts: Timestamp) -> &mut L7Summary {
if self.summaries.get_mut(key, ts).is_none() {
self.summaries.insert(key.clone(), L7Summary::default(), ts);
}
self.summaries
.get_mut(key, ts)
.expect("summary just inserted")
}
#[cfg(feature = "tls")]
pub fn observe_tls(&mut self, key: &K, hs: &crate::tls::TlsHandshake, ts: Timestamp) {
self.summary_mut(key, ts).observe_tls(hs);
}
#[cfg(feature = "http")]
pub fn observe_http(&mut self, key: &K, msg: &crate::http::HttpMessage, ts: Timestamp) {
self.summary_mut(key, ts).observe_http(msg);
}
#[cfg(feature = "dns")]
pub fn observe_dns_query(&mut self, key: &K, q: &crate::dns::DnsQuery, ts: Timestamp) {
self.summary_mut(key, ts).observe_dns_query(q);
}
#[cfg(feature = "dns")]
pub fn observe_dns_response(&mut self, key: &K, r: &crate::dns::DnsResponse, ts: Timestamp) {
self.summary_mut(key, ts).observe_dns_response(r);
}
pub fn evict_expired(&mut self, now: Timestamp) {
self.summaries.evict_expired(now);
}
pub fn forget(&mut self, key: &K) -> bool {
self.summaries.remove(key).is_some()
}
}
impl<K> FlowAnalyzer<K>
where
K: Hash + Eq + Clone + KeyFields,
{
pub fn finalize(&mut self, key: &K, stats: FlowStats) -> AnalyzedFlow<K> {
let l7 = self.summaries.remove(key).unwrap_or_default();
let ioc_hits = self.compute_ioc(key, &l7);
let risk = self.compute_risk(key, &l7, &ioc_hits);
AnalyzedFlow {
key: key.clone(),
stats,
l7,
risk,
ioc_hits,
}
}
pub fn snapshot(&self, key: &K, now: Timestamp, stats: FlowStats) -> AnalyzedFlow<K> {
let l7 = self.summaries.peek(key, now).cloned().unwrap_or_default();
let ioc_hits = self.compute_ioc(key, &l7);
let risk = self.compute_risk(key, &l7, &ioc_hits);
AnalyzedFlow {
key: key.clone(),
stats,
l7,
risk,
ioc_hits,
}
}
fn compute_ioc(&self, key: &K, l7: &L7Summary) -> Vec<IocMatch> {
let mut hits = Vec::new();
let Some(ioc) = &self.ioc else {
return hits;
};
if let Some(ip) = key.src_ip()
&& let Some(m) = ioc.contains_ip(ip)
{
hits.push(m);
}
if let Some(ip) = key.dest_ip()
&& let Some(m) = ioc.contains_ip(ip)
{
hits.push(m);
}
if let Some(sn) = &l7.server_name
&& let Some(m) = ioc.contains_domain(sn)
{
hits.push(m);
}
for q in &l7.dns_queries {
if let Some(m) = ioc.contains_domain(q) {
hits.push(m);
}
}
if let Some(j) = &l7.ja3
&& let Some(m) = ioc.contains(IocKind::Ja3, j)
{
hits.push(m);
}
if let Some(j) = &l7.ja4
&& let Some(m) = ioc.contains(IocKind::Ja4, j)
{
hits.push(m);
}
hits
}
fn compute_risk(&self, key: &K, l7: &L7Summary, ioc_hits: &[IocMatch]) -> FlowRisk {
let mut r = FlowRisk::empty();
if let Some(v) = l7.tls_version
&& v < 0x0303
{
r |= FlowRisk::TLS_OBSOLETE_VERSION;
}
if let Some(c) = l7.tls_cipher
&& is_weak_cipher(c)
{
r |= FlowRisk::TLS_WEAK_CIPHER;
}
for q in &l7.dns_queries {
r |= FlowRisk::from_dns(q, &self.dga);
}
if let Some(sn) = &l7.server_name {
r |= FlowRisk::from_dns(sn, &self.dga);
}
if let Some(detected) = l7.app_proto
&& let (Some(pid), Some(sp), Some(dp)) =
(key.protocol_identifier(), key.src_port(), key.dest_port())
{
let proto = l4_from_id(pid);
match crate::well_known::protocol_label(proto, sp, dp) {
Some(label) if label_matches(label, detected) => {}
Some(_) => r |= FlowRisk::PORT_PROTO_MISMATCH,
None => r |= FlowRisk::KNOWN_PROTO_NONSTD_PORT,
}
}
if ioc_hits
.iter()
.any(|m| matches!(m.kind, IocKind::Ja3 | IocKind::Ja4))
{
r |= FlowRisk::SUSPICIOUS_JA4;
}
r
}
}
fn l4_from_id(id: u8) -> crate::L4Proto {
match id {
6 => crate::L4Proto::Tcp,
17 => crate::L4Proto::Udp,
1 => crate::L4Proto::Icmp,
58 => crate::L4Proto::IcmpV6,
132 => crate::L4Proto::Sctp,
other => crate::L4Proto::Other(other),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::extract::FiveTupleKey;
use std::net::SocketAddr;
fn key(dport: u16) -> FiveTupleKey {
FiveTupleKey {
proto: crate::L4Proto::Tcp,
a: "10.0.0.9:40000".parse::<SocketAddr>().unwrap(),
b: format!("93.184.216.34:{dport}")
.parse::<SocketAddr>()
.unwrap(),
}
}
fn ttl() -> Duration {
Duration::from_secs(60)
}
#[test]
fn empty_flow_finalizes_clean() {
let mut a = FlowAnalyzer::<FiveTupleKey>::new(ttl());
let af = a.finalize(&key(443), FlowStats::default());
assert!(af.is_clean());
assert!(af.l7.is_empty());
assert_eq!(a.len(), 0);
}
#[cfg(feature = "tls")]
#[test]
fn tls_obsolete_and_weak_cipher_scored() {
use crate::tls::{TlsHandshake, TlsVersion};
let mut a = FlowAnalyzer::<FiveTupleKey>::new(ttl());
let k = key(443);
let hs = TlsHandshake {
sni: Some("example.com".into()),
version: Some(TlsVersion::Tls1_0),
cipher_suite: Some(0x0005), ..Default::default()
};
a.observe_tls(&k, &hs, Timestamp::default());
assert_eq!(a.len(), 1);
let af = a.finalize(&k, FlowStats::default());
assert!(af.risk.contains(FlowRisk::TLS_OBSOLETE_VERSION));
assert!(af.risk.contains(FlowRisk::TLS_WEAK_CIPHER));
assert!(!af.risk.contains(FlowRisk::PORT_PROTO_MISMATCH));
assert_eq!(af.l7.server_name.as_deref(), Some("example.com"));
assert_eq!(a.len(), 0);
}
#[cfg(feature = "tls")]
#[test]
fn ioc_ja4_hit_sets_suspicious_and_records_match() {
use crate::tls::TlsHandshake;
let mut ioc = IocSet::new();
ioc.insert(IocKind::Ja4, "t13d1516h2_abc_def", Some(95), Some("intel"));
let mut a = FlowAnalyzer::<FiveTupleKey>::new(ttl()).with_ioc(ioc);
let k = key(8443);
let hs = TlsHandshake {
ja4: Some("t13d1516h2_abc_def".into()),
..Default::default()
};
a.observe_tls(&k, &hs, Timestamp::default());
let af = a.finalize(&k, FlowStats::default());
assert!(af.has_ioc());
assert!(af.risk.contains(FlowRisk::SUSPICIOUS_JA4));
assert!(af.ioc_hits.iter().any(|m| m.kind == IocKind::Ja4));
}
#[cfg(feature = "dns")]
#[test]
fn dns_dga_and_ip_ioc() {
use crate::dns::{DnsFlags, DnsQuery, DnsQuestion};
let mut ioc = IocSet::new();
ioc.insert(IocKind::Ipv4, "93.184.216.34", Some(80), Some("c2-list"));
let mut a = FlowAnalyzer::<FiveTupleKey>::new(ttl()).with_ioc(ioc);
let k = key(53);
let q = DnsQuery {
transaction_id: 1,
flags: DnsFlags(0),
questions: vec![DnsQuestion {
name: "kq3v9z7xj2wq.com".into(),
qtype: 1,
qclass: 1,
}],
timestamp: Timestamp::default(),
};
a.observe_dns_query(&k, &q, Timestamp::default());
let af = a.finalize(&k, FlowStats::default());
assert!(af.risk.contains(FlowRisk::DGA_DOMAIN));
assert!(af.ioc_hits.iter().any(|m| m.kind == IocKind::Ipv4));
}
#[test]
fn evict_expired_reclaims_idle_state() {
let mut a = FlowAnalyzer::<FiveTupleKey>::with_capacity(Duration::from_secs(1), 100);
a.summaries
.insert(key(443), L7Summary::default(), Timestamp::new(0, 0));
assert_eq!(a.len(), 1);
a.evict_expired(Timestamp::new(5, 0));
assert_eq!(a.len(), 0);
}
}