use std::{ collections::HashMap, net::{Ipv4Addr, Ipv6Addr}, time::SystemTime, }; use crate::utils::ip::print_proto; use network_types::ip::{IpProto, Ipv4Hdr}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; #[derive(Serialize, Deserialize)] pub struct ConectionCoreAtom { pub len: u128, pub ip4_list: Vec>, pub ip6_list: Vec>, } #[derive(Serialize, Deserialize)] pub struct IpAtom { len: u128, ip_local: IpType, port_local: u16, ip_remote: IpType, port_remote: u16, //pub proto_len: HashMap pub proto_len: Vec<(u8, u128)>, } impl ConectionCoreAtom { pub fn new() -> Self { ConectionCoreAtom { len: 0, ip4_list: Vec::new(), ip6_list: Vec::new(), } } pub fn addv4( &mut self, source_ip: u32, dest_ip: u32, source_port: u16, dest_port: u16, protocol: IpProto, size: u16, ) { self.len += size as u128; let source_ip_format = Ipv4Addr::from_bits(source_ip); let dest_ip_format = Ipv4Addr::from_bits(dest_ip); let local_ip; let remote_ip; let local_port; let remote_port; if source_ip_format.is_private() { local_ip = source_ip_format; remote_ip = dest_ip_format; local_port = source_port; remote_port = dest_port; } else { local_ip = dest_ip_format; remote_ip = source_ip_format; local_port = dest_port; remote_port = dest_port; } match self.check_packagev4(local_ip, remote_ip, local_port, remote_port) { Some(ip_atom) => { ip_atom.add(protocol, size as u128); } None => { let mut ip_atom = IpAtom::new_ip(local_ip, remote_ip, local_port, remote_port); ip_atom.add(protocol, size as u128); self.ip4_list.push(ip_atom); } } } pub fn addv6( &mut self, source_ip: Ipv6Addr, dest_ip: Ipv6Addr, source_port: u16, dest_port: u16, protocol: IpProto, size: u16, ) { self.len += size as u128; let local_ip; let remote_ip; let local_port; let remote_port; if source_ip.is_unique_local() { local_ip = source_ip; local_port = source_port; remote_ip = dest_ip; remote_port = dest_port; } else { local_ip = dest_ip; local_port = dest_port; remote_ip = source_ip; remote_port = source_port; } match self.check_packagev6(local_ip, remote_ip, local_port, remote_port) { Some(ip_atom) => { ip_atom.add(protocol, size as u128); } None => { let mut ip_atom = IpAtom::new_ip(local_ip, remote_ip, local_port, remote_port); ip_atom.add(protocol, size as u128); self.ip6_list.push(ip_atom); } } } fn check_packagev4( &mut self, local_ip: Ipv4Addr, remote_ip: Ipv4Addr, local_port: u16, remote_port: u16, ) -> Option<&mut IpAtom> { let mut index = 0; loop { match self.ip4_list.get(index) { Some(ip) => { if ip.eq_splited(local_ip, remote_ip, local_port, remote_port) { return self.ip4_list.get_mut(index); } } None => { return None; } } index += 1; } } fn check_packagev6( &mut self, local_ip: Ipv6Addr, remote_ip: Ipv6Addr, local_port: u16, remote_port: u16, ) -> Option<&mut IpAtom> { let index = 0; loop { match self.ip6_list.get(index) { Some(ip) => { if ip.eq_splited(local_ip, remote_ip, local_port, remote_port) { return self.ip6_list.get_mut(index); } } None => return None, } } } pub fn generate_json(&self, timestamp: String) -> Value { json!({ "@timestamp" : timestamp, "ips_v4" : self.ip4_list, "ips_v6" : self.ip6_list, }) } pub fn reset(&mut self) -> Self { let old_len = self.len; self.len = 0; Self { len: old_len, ip4_list: self.ip4_list.drain(..).collect(), ip6_list: self.ip6_list.drain(..).collect(), } } /*fn create_ipv4(&mut self, source_ip: u32, dest_ip: u32) { let source_ip_format = Ipv4Addr::from_bits(source_ip); let dest_ip_format = Ipv4Addr::from_bits(dest_ip); self.ip4_list.push(if source_ip_format.is_private() { IpAtom::new_ip(source_ip_format, dest_ip_format) } else { IpAtom::new_ip(dest_ip_format, source_ip_format) }) } fn create_ipv6(&mut self, source_ip: Ipv6Addr, dest_ip: Ipv6Addr) { self.ip6_list.push(if source_ip.is_unique_local() { IpAtom::new_ip(source_ip, dest_ip) } else { IpAtom::new_ip(dest_ip, source_ip) }); }*/ } impl IpAtom { pub fn new_ip(local_ip: IpType, remote_ip: IpType, local_port: u16, remote_port: u16) -> Self { IpAtom { len: 0, ip_local: local_ip, port_local: local_port, ip_remote: remote_ip, port_remote: remote_port, proto_len: Vec::new(), } } //hashmap can be less efficient /*pub fn add(&mut self,protocol: IpProto, size: u128) { self.len+=size; match self.proto_len.get_mut(&(protocol as u8)) { Some(l) => { *l+=size; }, None => { self.proto_len.insert(protocol as u8,size); }, }; }*/ pub fn add(&mut self, protocol: IpProto, size: u128) { self.len += size; let mut index = 0; loop { match self.proto_len.get_mut(index) { Some(proto) => { if proto.0 == protocol as u8 { proto.1 += size; } } None => { self.proto_len.push((protocol as u8, size)); break; } } index += 1; } } pub fn eq_splited( &self, local_ip: IpType, remote_ip: IpType, local_port: u16, remote_port: u16, ) -> bool { ((self.ip_local == local_ip && self.port_local == local_port) && (self.ip_remote == remote_ip && self.port_remote == remote_port)) || ((self.ip_local == remote_ip && self.port_local == remote_port) && (self.ip_remote == local_ip && self.port_remote == local_port)) } }