net2elk/net-logger/src/clasificator/ip_group.rs

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<IpAtom<Ipv4Addr>>,
    pub ip6_list: Vec<IpAtom<Ipv6Addr>>,
}

#[derive(Serialize, Deserialize)]
pub struct IpAtom<IpType> {
    len: u128,
    ip_local: IpType,
    port_local: u16,
    ip_remote: IpType,
    port_remote: u16,
    //pub proto_len: HashMap<u8, u128>
    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<Ipv4Addr>> {
        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<Ipv6Addr>> {
        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<IpType: Eq> IpAtom<IpType> {
    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))
    }
}