// SPDX-FileCopyrightText: 2023 Greenbone AG
//
// SPDX-License-Identifier: GPL-2.0-or-later WITH x11vnc-openssl-exception

use std::{
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket},
    str::FromStr,
};

use crate::nasl::prelude::*;
use pcap::{Address, Device};
use pnet_base::MacAddr;

use super::{
    RawIpError,
    frame_forgery::{ArpFrame, ETHERTYPE_ARP, Frame},
};

use pnet::{
    datalink::interfaces,
    packet::{Packet, ipv4::Ipv4Packet, ipv6::Ipv6Packet, tcp::*, udp::MutableUdpPacket},
};
use socket2::{Domain, Protocol, Socket};

pub const DEFAULT_TTL: u8 = 255;
pub const IP_PPRTO_VERSION_IPV4: u8 = 4;
pub const IP_LENGTH: usize = 20;
pub const HEADER_LENGTH: u8 = 5;
pub const FIX_IPV6_HEADER_LENGTH: usize = 40;
pub const IPPROTO_IPV6: u8 = 6;

/// Return the MAC address, given the interface name
pub fn get_local_mac_address(name: &str) -> Result<MacAddr, FnError> {
    interfaces()
        .into_iter()
        .find(|x| x.name == *name)
        .and_then(|dev| dev.mac)
        .ok_or_else(|| RawIpError::FailedToGetLocalMacAddress.into())
}

/// Convert a string in a IpAddr
pub fn ipstr2ipaddr(ip_addr: &str) -> Result<IpAddr, FnError> {
    match IpAddr::from_str(ip_addr) {
        Ok(ip) => Ok(ip),
        Err(_) => Err(FnError::from(ArgumentError::WrongArgument(
            "Invalid IP address".to_string(),
        ))
        .with(ReturnValue(NaslValue::Null))),
    }
}

/// Tests whether a packet sent to IP is LIKELY to route through the
/// kernel localhost interface
pub fn islocalhost(addr: IpAddr) -> bool {
    // If it is not 0.0.0.0 or doesn't start with 127.0.0.1 then it
    // probably isn't localhost
    if !addr.is_loopback() || !addr.is_unspecified() {
        return false;
    }
    // It is not associated to a local interface.
    if let Err(_e) = get_interface_by_local_ip(addr) {
        return false;
    }

    true
}

/// Get the interface from the local ip
pub fn get_interface_by_local_ip(local_address: IpAddr) -> Result<Device, FnError> {
    // This fake IP is used for matching (and return false)
    // during the search of the interface in case an interface
    // doesn't have an associated address.

    let fake_ip = match local_address {
        IpAddr::V4(_) => IpAddr::V4(Ipv4Addr::UNSPECIFIED),
        IpAddr::V6(_) => IpAddr::V6(Ipv6Addr::UNSPECIFIED),
    };

    let fake_addr = Address {
        addr: fake_ip,
        broadcast_addr: None,
        netmask: None,
        dst_addr: None,
    };

    let ip_match = |ip: &Address| ip.addr.eq(&local_address);

    let devices = Device::list().map_err(|_| RawIpError::FailedToGetDeviceList)?;
    devices
        .into_iter()
        .find(|x| {
            local_address
                == (x.addresses.clone().into_iter().find(ip_match))
                    .unwrap_or_else(|| fake_addr.to_owned())
                    .addr
        })
        .ok_or(RawIpError::InvalidIpAddress.into())
}

pub fn get_mtu(target_ip: IpAddr) -> Result<usize, RawIpError> {
    let (_, mtu): (String, usize) =
        mtu::interface_and_mtu(target_ip).map_err(|_| RawIpError::FailedToGetDeviceMTU)?;
    Ok(mtu)
}

fn bind_local_socket(dst: &SocketAddr) -> Result<UdpSocket, RawIpError> {
    match dst {
        SocketAddr::V4(_) => UdpSocket::bind("0.0.0.0:0"),
        SocketAddr::V6(_) => UdpSocket::bind("[0:0:0:0:0:0:0:0]:0"),
    }
    .map_err(RawIpError::FailedToBind)
}

/// Return the source IP address given the destination IP address
pub fn get_source_ip(dst: IpAddr) -> Result<IpAddr, FnError> {
    let fake_port = 50000u16;
    let socket = SocketAddr::new(dst, fake_port);
    let sd = format!("{dst}:{fake_port}");
    let local_socket = bind_local_socket(&socket)?;
    local_socket
        .connect(sd)
        .ok()
        .and_then(|_| local_socket.local_addr().ok())
        .and_then(|l_addr| IpAddr::from_str(&l_addr.ip().to_string()).ok())
        .ok_or(RawIpError::NoRouteToDestination.into())
}

pub fn get_source_ipv6(dst: Ipv6Addr) -> Result<Ipv6Addr, FnError> {
    match get_source_ip(IpAddr::from(dst)) {
        Ok(IpAddr::V6(a)) => Ok(a),
        _ => Err(RawIpError::NoRouteToDestination.into()),
    }
}

pub fn get_source_ipv4(dst: Ipv4Addr) -> Result<Ipv4Addr, FnError> {
    match get_source_ip(IpAddr::from(dst)) {
        Ok(IpAddr::V4(a)) => Ok(a),
        _ => Err(RawIpError::NoRouteToDestination.into()),
    }
}

pub trait ChecksumCalculator<'a, V: 'a> {
    fn calculate_checksum(&self, chksum: Option<u16>, pkt: &'a V) -> u16;
}

impl<'a> ChecksumCalculator<'a, Ipv4Packet<'a>> for MutableUdpPacket<'a> {
    fn calculate_checksum(&self, chksum: Option<u16>, pkt: &'a Ipv4Packet) -> u16 {
        let chksum = chksum.unwrap_or(0);
        if chksum != 0 {
            return chksum.to_be();
        }
        pnet::packet::udp::ipv4_checksum(
            &self.to_immutable(),
            &pkt.get_source(),
            &pkt.get_destination(),
        )
    }
}

impl<'a> ChecksumCalculator<'a, Ipv4Packet<'a>> for MutableTcpPacket<'a> {
    fn calculate_checksum(&self, chksum: Option<u16>, pkt: &'a Ipv4Packet) -> u16 {
        let chksum = chksum.unwrap_or(0);
        if chksum != 0 {
            return chksum.to_be();
        }
        pnet::packet::tcp::ipv4_checksum(
            &self.to_immutable(),
            &pkt.get_source(),
            &pkt.get_destination(),
        )
    }
}

impl<'a> ChecksumCalculator<'a, Ipv6Packet<'a>> for MutableUdpPacket<'a> {
    fn calculate_checksum(&self, chksum: Option<u16>, pkt: &'a Ipv6Packet) -> u16 {
        let chksum = chksum.unwrap_or(0);
        if chksum != 0 {
            return chksum.to_be();
        }
        pnet::packet::udp::ipv6_checksum(
            &self.to_immutable(),
            &pkt.get_source(),
            &pkt.get_destination(),
        )
    }
}

impl<'a> ChecksumCalculator<'a, Ipv6Packet<'a>> for MutableTcpPacket<'a> {
    fn calculate_checksum(&self, chksum: Option<u16>, pkt: &'a Ipv6Packet) -> u16 {
        let chksum = chksum.unwrap_or(0);
        if chksum != 0 {
            return chksum.to_be();
        }
        pnet::packet::tcp::ipv6_checksum(
            &self.to_immutable(),
            &pkt.get_source(),
            &pkt.get_destination(),
        )
    }
}

/// Forge a data link layer frame with an ARP request in the payload
pub fn forge_arp_frame(eth_src: MacAddr, src_ip: Ipv4Addr, dst_ip: Ipv4Addr) -> Vec<u8> {
    let mut frame = Frame::new();
    frame.set_srchaddr(eth_src);
    frame.set_dsthaddr(MacAddr::broadcast());
    frame.set_ethertype(ETHERTYPE_ARP.to_le());

    let mut arp_frame = ArpFrame::new();
    arp_frame.set_srchaddr(eth_src);
    arp_frame.set_srcip(src_ip);
    arp_frame.set_dsthaddr(MacAddr::zero());
    arp_frame.set_dstip(dst_ip);

    frame.set_payload(arp_frame.into());
    frame.into()
}

const IPPROTO_RAW: i32 = 255;

fn new_raw_socket_v4() -> Result<Socket, RawIpError> {
    Socket::new_raw(
        Domain::IPV4,
        socket2::Type::RAW,
        Some(Protocol::from(IPPROTO_RAW)),
    )
    .map_err(RawIpError::FailedToBind)
}

fn new_raw_socket_v6() -> Result<Socket, RawIpError> {
    Socket::new_raw(
        Domain::IPV6,
        socket2::Type::RAW,
        Some(Protocol::from(IPPROTO_RAW)),
    )
    .map_err(RawIpError::FailedToBind)
}

// Send ipv6 packet
pub fn send_v6_packet(pkt: Ipv6Packet<'static>) -> Result<(), RawIpError> {
    tracing::debug!("starting sending packet");
    let sock = new_raw_socket_v6()?;
    sock.set_header_included_v6(true)
        .map_err(RawIpError::FailedToBind)?;

    let sockaddr = SocketAddr::new(IpAddr::from(pkt.get_destination()), 0);
    match sock.send_to(pkt.packet(), &sockaddr.into()) {
        Ok(b) => {
            tracing::debug!("Sent {} bytes", b);
        }
        Err(e) => {
            return Err(RawIpError::SendPacket(e.to_string()));
        }
    };
    Ok(())
}

// Send ipv4 packet
pub fn send_v4_packet(pkt: Ipv4Packet<'static>) -> Result<(), RawIpError> {
    tracing::debug!("starting sending packet");
    let sock = new_raw_socket_v4()?;
    sock.set_header_included_v4(true)
        .map_err(RawIpError::FailedToBind)?;

    let sockaddr = SocketAddr::new(IpAddr::from(pkt.get_destination()), 0);
    match sock.send_to(pkt.packet(), &sockaddr.into()) {
        Ok(b) => {
            tracing::debug!("Sent {} bytes", b);
        }
        Err(e) => {
            return Err(RawIpError::SendPacket(e.to_string()));
        }
    };
    Ok(())
}
