bgp-packet 0.0.2

use std::net::Ipv4Addr;

use bitfield::bitfield;
use bytes::{BufMut, BytesMut};
use eyre::{Result, bail, eyre};
use nom::{
    Err::Failure,
    IResult, Parser,
    combinator::peek,
    number::complete::{be_u8, be_u16, be_u24, be_u32},
};
use serde::{Deserialize, Serialize};
use strum::EnumDiscriminants;

use crate::{
    constants::{AddressFamilyId, SubsequentAfi},
    parser::{BgpParserError, ParserContext, ProtocolErrorValues},
};

macro_rules! fail {
    ( $($arg:expr),* $(,)? ) => {
        Failure(BgpParserError::Eyre(eyre!($($arg),*)))
    };
}

/// Contains messages related to the BGP Open message,
/// i.e. the Open message itself and options / capabilities.
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct OpenMessage {
    /// Version of the BGP protocol in use.
    pub version: u8,
    /// AS Number of the BGP speaker, or AS_TRANS if using 4 byte ASN.
    pub asn: u16,
    /// Hold time parameter in seconds.
    pub hold_time: u16,
    /// Global identifier of the BGP speaker.
    pub identifier: Ipv4Addr,
    /// Options.
    pub options: Vec<OpenOption>,
}

impl OpenMessage {
    /// The from_wire implementation parses a BGP Open message from a buffer
    /// containing the payload of the open message (i.e. without the BGP header / type / length)
    /// and the buf passed in must contain only the open message.
    pub fn from_wire<'a>(
        ctx: &ParserContext,
        buf: &'a [u8],
    ) -> IResult<&'a [u8], Self, BgpParserError<&'a [u8]>> {
        let (buf, version) = be_u8(buf)?;
        if version != 4 {
            return Err(Failure(BgpParserError::ProtocolError(
                ProtocolErrorValues::UnknownOpenOption(version),
            )));
        }

        let (buf, asn) = be_u16(buf)?;
        let (buf, hold_time) = be_u16(buf)?;
        let (buf, identifier) = be_u32(buf)?;

        // RFC 9072: Extended Open Options length.
        // To encode a capabilities length greater than 255 the packet structure
        // is prefixed with the extended length option (255), a length of 255,
        // and then a u16 containing the payload option.
        let (buf, opt_len) = be_u8(buf)?;
        let (_, inner) = peek(be_u8).parse(buf)?;

        let (buf, opt_len) = if inner == 255 {
            let (buf, _) = be_u8(buf)?; // Skip past the inner type.
            be_u16(buf)?
        } else {
            (buf, opt_len as u16)
        };

        let (buf, opt_buf) = nom::bytes::take(opt_len).parse(buf)?;
        let (opt_buf, options) =
            nom::multi::many0(|buf| OpenOption::from_wire(ctx, buf)).parse(opt_buf)?;

        if !opt_buf.is_empty() {
            return Err(fail!(
                "Leftover bytes after parsing open options: {:x?}",
                opt_buf
            ));
        }

        Ok((
            buf,
            Self {
                version,
                asn,
                hold_time,
                identifier: Ipv4Addr::from(identifier),
                options,
            },
        ))
    }

    pub fn to_wire(&self, ctx: &ParserContext, out: &mut BytesMut) -> Result<()> {
        out.put_u8(self.version);
        out.put_u16(self.asn);
        out.put_u16(self.hold_time);
        out.put_u32(self.identifier.into());

        // RFC 9072: Depending on whether the options fit in < 255 bytes
        // we encode the length accordingly.
        let opt_len = self
            .options
            .iter()
            .map(|o| o.wire_len(ctx))
            .sum::<Result<u16>>()?;
        if opt_len > u8::MAX as u16 {
            out.put_u8(0xff); // Set length to 255 as per RFC 9072.
            out.put_u8(0xff); // Open option type for Ext length.
            out.put_u16(opt_len);
        } else {
            out.put_u8(opt_len as u8);
        }

        for option in &self.options {
            option.to_wire(ctx, out)?;
        }

        Ok(())
    }

    pub fn wire_len(&self, ctx: &ParserContext) -> Result<u16> {
        let mut total = 10;
        let opt_len = self
            .options
            .iter()
            .map(|o| o.wire_len(ctx))
            .sum::<Result<u16>>()?;
        if opt_len > u8::MAX as u16 {
            total += opt_len + 3; // 1+2 bytes extra for the type + extended length.
        } else {
            total += opt_len;
        }
        Ok(total)
    }
}

#[derive(Debug, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum OpenOption {
    Capabilities(Vec<Capability>) = 2,
}

impl OpenOption {
    pub fn from_wire<'a>(
        ctx: &ParserContext,
        buf: &'a [u8],
    ) -> IResult<&'a [u8], Self, BgpParserError<&'a [u8]>> {
        let (buf, type_code) = be_u8(buf)?;
        let (buf, length) = be_u8(buf)?;

        let (buf, payload) = nom::bytes::take(length).parse(buf)?;

        match type_code {
            v if v == OpenOptionDiscriminants::Capabilities as u8 => {
                let (payload_consumed, (capabilities, _eof)) = nom::multi::many_till(
                    |buf| Capability::from_wire(ctx, buf),
                    nom::combinator::eof,
                )
                .parse(payload)?;

                if !payload_consumed.is_empty() {
                    return Err(fail!(
                        "Leftover bytes when parsing BGP capabilities: {:x?}",
                        payload_consumed
                    ));
                }
                Ok((buf, Self::Capabilities(capabilities)))
            }
            other => {
                return Err(Failure(BgpParserError::ProtocolError(
                    ProtocolErrorValues::UnknownOpenOption(other),
                )));
            }
        }
    }

    pub fn to_wire(&self, ctx: &ParserContext, out: &mut BytesMut) -> Result<()> {
        match self {
            OpenOption::Capabilities(items) => {
                let cap_len = items.iter().map(|i| i.wire_len(ctx)).sum::<Result<u16>>()?;
                if cap_len > u8::MAX as u16 {
                    bail!("Cannot write Capabilities with more than u8::MAX length");
                }
                out.put_u8(OpenOptionDiscriminants::Capabilities as u8);
                out.put_u8(cap_len as u8);
                for item in items {
                    item.to_wire(ctx, out)?;
                }
            }
        }
        Ok(())
    }

    pub fn wire_len(&self, ctx: &ParserContext) -> Result<u16> {
        match self {
            OpenOption::Capabilities(items) => items
                .iter()
                .map(|i| i.wire_len(ctx))
                .sum::<Result<u16>>()
                .map(|v| v + 2),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum Capability {
    /// MultiProtocol Extension (RFC 2858).
    MultiProtocol {
        afi: AddressFamilyId,
        safi: SubsequentAfi,
    } = 1,
    /// Route Refresh capability (RFC 2918).
    RouteRefresh = 2,
    /// Outbound Route Filtering (RFC 5291).
    /// https://datatracker.ietf.org/doc/html/rfc5291
    OutboundRouteFilter {
        afi: AddressFamilyId,
        safi: SubsequentAfi,
        orfs: Vec<OutboundRouteFilterEntry>,
    } = 3,
    /// Extended Next Hop encoding (RFC 8950).
    ExtendedNextHop {
        entries: Vec<ExtendedNextHop>,
    } = 5,
    /// Extended Message (RFC 8654).
    ExtendedMessage = 6,
    /// BGPSec (RFC 8205).
    BgpSec {
        header: BgpSecHeader,
        afi: AddressFamilyId,
    } = 7,
    /// Multiple labels compatibility (RFC 8277).
    MultiLabelCompat {} = 8,
    BgpRole {
        role: BgpRole,
    } = 9,
    /// Graceful restart capability (RFC 4724).
    GracefulRestart {
        /// When set to true it means that the BGP speaker has just restarted
        /// and the peer must not wait for the End of RIB marker from the
        /// speaker before advertising routing information to the speaker.
        restart_state: bool,
        /// This is the estimated time (in seconds) it will take for the
        /// BGP session to be re-established after a restart.  This can be
        /// used to speed up routing convergence by its peer in case that
        /// the BGP speaker does not come back after a restart.
        restart_time: u16,
        /// The AFI and SAFI, taken in combination, indicate that Graceful
        /// Restart is supported for routes that are advertised with the
        /// same AFI and SAFI.  Routes may be explicitly associated with a
        /// particular AFI and SAFI using the encoding of [BGP-MP] or
        /// implicitly associated with <AFI=IPv4, SAFI=Unicast> if using
        /// the encoding of [BGP-4].
        entries: Vec<GracefulRestartEntry>,
    } = 64,
    /// Four Byte ASN (RFC 4274).
    FourByteAsn {
        asn: u32,
    } = 65,
    /// Additional Path (RFC 7911).
    AddPath {
        afi: AddressFamilyId,
        safi: SubsequentAfi,
        send_recv: AddPathSendRecv,
    } = 69,
    /// Enhanced Route Refresh (RFC 7313).
    EnhancedRouteRefresh = 70,
    /// Long Lived Graceful Restart (RFC 9494).
    LongLivedGracefulRestart(Vec<LongLivedGracefulRestart>) = 71,
    /// Unknown Capability encapsulates any capabilities that we don't know about.
    UnknownCapability {
        type_code: u8,
        length: u8,
        value: Vec<u8>,
    },
}

impl Capability {
    fn discriminant(&self) -> u8 {
        match self {
            Capability::MultiProtocol { .. } => CapabilityDiscriminants::MultiProtocol as u8,
            Capability::RouteRefresh => CapabilityDiscriminants::RouteRefresh as u8,
            Capability::OutboundRouteFilter { .. } => {
                CapabilityDiscriminants::OutboundRouteFilter as u8
            }
            Capability::ExtendedNextHop { .. } => CapabilityDiscriminants::ExtendedNextHop as u8,
            Capability::ExtendedMessage => CapabilityDiscriminants::ExtendedMessage as u8,
            Capability::BgpSec { .. } => CapabilityDiscriminants::BgpSec as u8,
            Capability::MultiLabelCompat {} => CapabilityDiscriminants::MultiLabelCompat as u8,
            Capability::BgpRole { .. } => CapabilityDiscriminants::BgpRole as u8,
            Capability::GracefulRestart { .. } => CapabilityDiscriminants::GracefulRestart as u8,
            Capability::FourByteAsn { .. } => CapabilityDiscriminants::FourByteAsn as u8,
            Capability::AddPath { .. } => CapabilityDiscriminants::AddPath as u8,
            Capability::EnhancedRouteRefresh => CapabilityDiscriminants::EnhancedRouteRefresh as u8,
            Capability::LongLivedGracefulRestart(_) => {
                CapabilityDiscriminants::LongLivedGracefulRestart as u8
            }
            Capability::UnknownCapability { type_code, .. } => *type_code,
        }
    }
    /// `from_wire` parses a Capability message.
    pub fn from_wire<'a>(
        _ctx: &ParserContext,
        buf: &'a [u8],
    ) -> IResult<&'a [u8], Self, BgpParserError<&'a [u8]>> {
        let (buf, typ) = be_u8(buf)?;
        let (buf, len) = be_u8(buf)?;
        let (buf, cap_buf) = nom::bytes::take(len).parse(buf)?;
        macro_rules! require_empty {
            ($buf:expr, $cap_type:expr) => {
                if !buf.is_empty() {
                    return Err(nom::Err::Failure(BgpParserError::Eyre(eyre!(
                        "Unexpected leftover bytes {:x?} while parsing {} capability",
                        $buf,
                        $cap_type
                    ))));
                }
            };
        }
        match typ {
            v if v == CapabilityDiscriminants::MultiProtocol as u8 => {
                // MultiProtocol capability consists of an (AFI, SAFI).
                let (cap_buf, afi) = be_u16(cap_buf)?;
                let (cap_buf, _) = be_u8(cap_buf)?;
                let (cap_buf, safi) = be_u8(cap_buf)?;
                require_empty!(cap_buf, "MultiProtocol");
                let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;
                Ok((buf, Self::MultiProtocol { afi, safi }))
            }
            v if v == CapabilityDiscriminants::RouteRefresh as u8 => {
                require_empty!(cap_buf, "RouteRefresh");
                Ok((buf, Self::RouteRefresh))
            }
            v if v == CapabilityDiscriminants::OutboundRouteFilter as u8 => {
                let (cap_buf, afi) = be_u16(cap_buf)?;
                let (cap_buf, _) = be_u8(cap_buf)?; // Move past reverved byte.
                let (cap_buf, safi) = be_u8(cap_buf)?;
                let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;

                let (cap_buf, _num_orfs) = be_u8(cap_buf)?;

                let parse_orf = |buf: &'a [u8]| -> IResult<
                    &'a [u8],
                    OutboundRouteFilterEntry,
                    BgpParserError<&'a [u8]>,
                > {
                    let (buf, orf_type) = be_u8(buf)?;
                    let (buf, send_recv) = be_u8(buf)?;
                    let orf_type = OrfType::from(orf_type);
                    let send_recv = OrfSendRecv::try_from(send_recv)
                        .map_err(|e| fail!("Failed to parse ORF send/recv: {}", e))?;

                    Ok((
                        buf,
                        OutboundRouteFilterEntry {
                            r#type: orf_type,
                            send_recv,
                        },
                    ))
                };

                let (cap_buf, orfs) = nom::multi::many0(parse_orf).parse(cap_buf)?;
                require_empty!(cap_buf, "OutboundRouteFilter");

                Ok((buf, Self::OutboundRouteFilter { afi, safi, orfs }))
            }
            v if v == CapabilityDiscriminants::ExtendedNextHop as u8 => {
                // parse_enh parses a single extended nexthop entry.
                let parse_enh = |buf: &'a [u8]| -> IResult<
                    &'a [u8],
                    ExtendedNextHop,
                    BgpParserError<&'a [u8]>,
                > {
                    let (buf, afi) = be_u16(buf)?;
                    let (buf, safi) = be_u16(buf)?;
                    let (buf, nh_afi) = be_u16(buf)?;
                    let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                    let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;
                    let nh_afi = AddressFamilyId::try_from(nh_afi).map_err(|e| fail!("{}", e))?;
                    Ok((buf, ExtendedNextHop { afi, safi, nh_afi }))
                };
                let (cap_buf, entries) = nom::multi::many1(parse_enh).parse(cap_buf)?;
                require_empty!(cap_buf, "ExtendedNextHop");

                Ok((cap_buf, Self::ExtendedNextHop { entries }))
            }
            v if v == CapabilityDiscriminants::ExtendedMessage as u8 => {
                require_empty!(cap_buf, "ExtendedMessage");
                Ok((cap_buf, Self::ExtendedMessage))
            }
            v if v == CapabilityDiscriminants::BgpSec as u8 => {
                let (cap_buf, header) = be_u8(cap_buf)?;
                let (cap_buf, afi) = be_u16(cap_buf)?;
                let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                require_empty!(cap_buf, "BgpSec");
                Ok((
                    cap_buf,
                    Self::BgpSec {
                        header: BgpSecHeader(header),
                        afi,
                    },
                ))
            }
            v if v == CapabilityDiscriminants::MultiLabelCompat as u8 => {
                todo!()
            }
            v if v == CapabilityDiscriminants::BgpRole as u8 => {
                let (cap_buf, role) = be_u8(cap_buf)?;
                let role = BgpRole::try_from(role).map_err(|e| fail!("{}", e))?;
                require_empty!(cap_buf, "BgpRole");
                Ok((cap_buf, Self::BgpRole { role }))
            }
            v if v == CapabilityDiscriminants::GracefulRestart as u8 => {
                // flags_restart_time is: (4 bits: restart flags, 12 bits: restart time (s)).
                let (cap_buf, flags_restart_time) = be_u16(cap_buf)?;
                let restart_state = ((1 << 15) & flags_restart_time) != 0;
                let restart_time = flags_restart_time & 0xfff;
                let parse_entry = |buf: &'a [u8]| -> IResult<
                    &'a [u8],
                    GracefulRestartEntry,
                    BgpParserError<&'a [u8]>,
                > {
                    let (buf, afi) = be_u16(buf)?;
                    let (buf, safi) = be_u8(buf)?;
                    let (buf, flags) = be_u8(buf)?;
                    let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                    let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;
                    Ok((
                        buf,
                        GracefulRestartEntry {
                            afi,
                            safi,
                            preserve_forwarding: (flags & (1 << 7)) != 0,
                        },
                    ))
                };
                let (cap_buf, entries) = nom::multi::many0(parse_entry).parse(cap_buf)?;
                require_empty!(cap_buf, "GracefulRestart");
                Ok((
                    cap_buf,
                    Self::GracefulRestart {
                        restart_state,
                        restart_time,
                        entries,
                    },
                ))
            }
            v if v == CapabilityDiscriminants::FourByteAsn as u8 => {
                let (cap_buf, asn) = be_u32(cap_buf)?;
                require_empty!(cap_buf, "FourByteAsn");
                Ok((cap_buf, Self::FourByteAsn { asn }))
            }
            v if v == CapabilityDiscriminants::AddPath as u8 => {
                let (cap_buf, afi) = be_u16(cap_buf)?;
                let (cap_buf, safi) = be_u8(cap_buf)?;
                let (cap_buf, send_recv) = be_u8(cap_buf)?;
                let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;
                let send_recv = AddPathSendRecv::try_from(send_recv).map_err(|e| fail!("{}", e))?;
                require_empty!(cap_buf, "AddPath");
                Ok((
                    cap_buf,
                    Self::AddPath {
                        afi,
                        safi,
                        send_recv,
                    },
                ))
            }
            v if v == CapabilityDiscriminants::EnhancedRouteRefresh as u8 => {
                require_empty!(cap_buf, "EnhancedRouteRefresh");
                Ok((cap_buf, Self::EnhancedRouteRefresh))
            }
            v if v == CapabilityDiscriminants::LongLivedGracefulRestart as u8 => {
                let parse_llgr = |buf: &'a [u8]| -> IResult<
                    &'a [u8],
                    LongLivedGracefulRestart,
                    BgpParserError<&'a [u8]>,
                > {
                    let (buf, afi) = be_u16(buf)?;
                    let (buf, safi) = be_u8(buf)?;
                    let (buf, flags) = be_u8(buf)?;
                    let (buf, stale_time) = be_u24(buf)?;
                    let afi = AddressFamilyId::try_from(afi).map_err(|e| fail!("{}", e))?;
                    let safi = SubsequentAfi::try_from(safi).map_err(|e| fail!("{}", e))?;
                    Ok((
                        buf,
                        LongLivedGracefulRestart {
                            afi,
                            safi,
                            flags,
                            stale_time,
                        },
                    ))
                };
                let (cap_buf, entries) = nom::multi::many0(parse_llgr).parse(cap_buf)?;
                require_empty!(cap_buf, "LongLivedGracefulRestart");
                Ok((cap_buf, Self::LongLivedGracefulRestart(entries)))
            }
            type_code => {
                // Handling an unknown capability, so just parse it out and store the value
                // as bytes.
                Ok((
                    &[],
                    Self::UnknownCapability {
                        type_code,
                        length: cap_buf.len() as u8,
                        value: cap_buf.to_vec(),
                    },
                ))
            }
        }
    }

    pub fn to_wire(&self, _ctx: &ParserContext, out: &mut BytesMut) -> Result<()> {
        out.put_u8(self.discriminant());
        out.put_u8(self.payload_len());
        match self {
            Capability::MultiProtocol { afi, safi } => {
                out.put_u16(*afi as u16);
                out.put_u8(0); // Reserved byte.
                out.put_u8(*safi as u8);
            }
            Capability::RouteRefresh => {}
            Capability::OutboundRouteFilter { afi, safi, orfs } => {
                out.put_u16(*afi as u16);
                out.put_u8(0); // Reserved byte.
                out.put_u8(*safi as u8);
                for orf in orfs {
                    out.put_u8(orf.r#type.clone().into());
                    out.put_u8(orf.send_recv.clone() as u8);
                }
            }
            Capability::ExtendedNextHop { entries } => {
                for entry in entries {
                    out.put_u16(entry.afi as u16);
                    out.put_u16(entry.safi as u16);
                    out.put_u16(entry.nh_afi as u16);
                }
            }
            Capability::ExtendedMessage => {}
            Capability::BgpSec { header, afi } => {
                out.put_u8(header.0);
                out.put_u16(*afi as u16);
            }
            Capability::MultiLabelCompat {} => todo!(),
            Capability::BgpRole { role } => {
                out.put_u8(role.into());
            }
            Capability::GracefulRestart {
                restart_state,
                restart_time,
                entries,
            } => {
                let mut flags_restart_time = 0;
                if *restart_state {
                    flags_restart_time |= 1 << 15;
                }
                if *restart_time > 0xfff {
                    bail!("Restart time too large to be encoded: {}", restart_time);
                }
                flags_restart_time |= restart_time;
                out.put_u16(flags_restart_time);
                for entry in entries {
                    out.put_u16(entry.afi as u16);
                    out.put_u8(entry.safi as u8);
                    let mut flags = 0;
                    if entry.preserve_forwarding {
                        flags |= 1 << 7;
                    }
                    out.put_u8(flags);
                }
            }
            Capability::FourByteAsn { asn } => {
                out.put_u32(*asn);
            }
            Capability::AddPath {
                afi,
                safi,
                send_recv,
            } => {
                out.put_u16(*afi as u16);
                out.put_u8(*safi as u8);
                out.put_u8(send_recv.into());
            }
            Capability::EnhancedRouteRefresh => {}
            Capability::LongLivedGracefulRestart(llgrs) => {
                for llgr in llgrs {
                    out.put_u16(llgr.afi as u16);
                    out.put_u8(llgr.safi as u8);
                    out.put_u8(llgr.flags as u8);
                    out.put(&llgr.stale_time.to_be_bytes()[1..]);
                }
            }
            Capability::UnknownCapability { value, .. } => out.put(value.as_slice()),
        }
        Ok(())
    }

    pub fn wire_len(&self, _ctx: &ParserContext) -> Result<u16> {
        Ok(2 + self.payload_len() as u16)
    }

    fn payload_len(&self) -> u8 {
        match self {
            Capability::MultiProtocol { .. } => 4,
            Capability::RouteRefresh => 0,
            Capability::OutboundRouteFilter { orfs, .. } => 4 + 2 * orfs.len() as u8,
            Capability::ExtendedNextHop { entries } => 6 * entries.len() as u8,
            Capability::ExtendedMessage => 0,
            Capability::BgpSec { .. } => 3,
            Capability::MultiLabelCompat {} => todo!(),
            Capability::BgpRole { .. } => 1,
            Capability::GracefulRestart { entries, .. } => 2 + 4 * entries.len() as u8,
            Capability::FourByteAsn { .. } => 4,
            Capability::AddPath { .. } => 4,
            Capability::EnhancedRouteRefresh => 0,
            Capability::LongLivedGracefulRestart(long_lived_graceful_restarts) => {
                8 * long_lived_graceful_restarts.len() as u8
            }
            Capability::UnknownCapability { value, .. } => value.len() as u8,
        }
    }
}

/// Represents an entry for the OutboundRouteFilter capability.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OutboundRouteFilterEntry {
    pub r#type: OrfType,
    pub send_recv: OrfSendRecv,
}

#[derive(Debug, Clone, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum OrfType {
    Reserved = 0,
    Unassigned(u8),
    AddressPrefix = 64,
    CoveringPrefix = 65,
    VpnPrefix = 66,
    VendorSpecific(u8),
}

impl From<u8> for OrfType {
    fn from(value: u8) -> Self {
        match value {
            v if v == OrfTypeDiscriminants::Reserved as u8 => Self::Reserved,
            v if (1..64).contains(&v) => Self::Unassigned(v),
            v if v == OrfTypeDiscriminants::AddressPrefix as u8 => Self::AddressPrefix,
            v if v == OrfTypeDiscriminants::CoveringPrefix as u8 => Self::CoveringPrefix,
            v if v == OrfTypeDiscriminants::VpnPrefix as u8 => Self::VpnPrefix,
            other => Self::VendorSpecific(other),
        }
    }
}

impl From<OrfType> for u8 {
    fn from(value: OrfType) -> Self {
        match value {
            OrfType::Reserved => OrfTypeDiscriminants::Reserved as u8,
            OrfType::Unassigned(inner) => inner,
            OrfType::AddressPrefix => OrfTypeDiscriminants::AddressPrefix as u8,
            OrfType::CoveringPrefix => OrfTypeDiscriminants::CoveringPrefix as u8,
            OrfType::VpnPrefix => OrfTypeDiscriminants::VpnPrefix as u8,
            OrfType::VendorSpecific(inner) => inner,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum OrfSendRecv {
    Receive = 1,
    Send = 2,
    Both = 3,
}

impl TryFrom<u8> for OrfSendRecv {
    type Error = eyre::ErrReport;

    fn try_from(value: u8) -> Result<Self> {
        Ok(match value {
            value if value == OrfSendRecv::Receive as u8 => OrfSendRecv::Receive,
            value if value == OrfSendRecv::Send as u8 => OrfSendRecv::Send,
            value if value == OrfSendRecv::Both as u8 => OrfSendRecv::Both,
            other => bail!("Invalid ORF send/recv value: {}", other),
        })
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum AddPathSendRecv {
    Receive = 1,
    Send = 2,
    Both = 3,
}

impl TryFrom<u8> for AddPathSendRecv {
    type Error = eyre::ErrReport;

    fn try_from(value: u8) -> Result<Self> {
        Ok(match value {
            value if value == AddPathSendRecv::Receive as u8 => AddPathSendRecv::Receive,
            value if value == AddPathSendRecv::Send as u8 => AddPathSendRecv::Send,
            value if value == AddPathSendRecv::Both as u8 => AddPathSendRecv::Both,
            other => bail!("Invalid AddPathSendRecv send/recv value: {}", other),
        })
    }
}

impl From<&AddPathSendRecv> for u8 {
    fn from(value: &AddPathSendRecv) -> Self {
        match value {
            AddPathSendRecv::Receive => AddPathSendRecv::Receive as u8,
            AddPathSendRecv::Send => AddPathSendRecv::Send as u8,
            AddPathSendRecv::Both => AddPathSendRecv::Both as u8,
        }
    }
}

/// `ExtendedNextHop` allows for the nexthop to be for a different address family
/// than the one being used to advertise prefixes for. E.g. it allows for sending
/// IPv4 prefixes with an IPv6 nexthop.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExtendedNextHop {
    /// The Address Family to update the nexthop format for.
    pub afi: AddressFamilyId,
    /// The Subsequent Address Family to update the nexthop format for.
    pub safi: SubsequentAfi,
    /// The Address Family that nexthops will be encoded as.
    pub nh_afi: AddressFamilyId,
}

#[derive(Debug, Clone, Serialize, Deserialize, EnumDiscriminants, PartialEq, Eq)]
#[repr(u8)]
pub enum BgpRole {
    Provider = 0,
    RouteServer = 1,
    RouteServerClient = 2,
    Customer = 3,
    Peer = 4,
    Unknown(u8),
}

impl From<u8> for BgpRole {
    fn from(value: u8) -> Self {
        match value {
            v if v == BgpRoleDiscriminants::Provider as u8 => Self::Provider,
            v if v == BgpRoleDiscriminants::RouteServer as u8 => Self::RouteServer,
            v if v == BgpRoleDiscriminants::RouteServerClient as u8 => Self::RouteServerClient,
            v if v == BgpRoleDiscriminants::Customer as u8 => Self::Customer,
            v if v == BgpRoleDiscriminants::Peer as u8 => Self::Peer,
            other => Self::Unknown(other),
        }
    }
}

impl From<&BgpRole> for u8 {
    fn from(value: &BgpRole) -> Self {
        match value {
            BgpRole::Provider => BgpRoleDiscriminants::Provider as u8,
            BgpRole::RouteServer => BgpRoleDiscriminants::RouteServer as u8,
            BgpRole::RouteServerClient => BgpRoleDiscriminants::RouteServerClient as u8,
            BgpRole::Customer => BgpRoleDiscriminants::Customer as u8,
            BgpRole::Peer => BgpRoleDiscriminants::Peer as u8,
            BgpRole::Unknown(other) => *other,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct LongLivedGracefulRestart {
    pub afi: AddressFamilyId,
    pub safi: SubsequentAfi,
    pub flags: u8,
    /// stale_time is the time in seconds after which routes are considered stale.
    /// Note that this is represented as a u24 on the wire.
    pub stale_time: u32,
}

bitfield! {
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
    pub struct BgpSecHeader(u8);
    impl new;
    u8,   version, set_version   : 3, 0;   // 4 bits
    bool, dir,     set_dir       : 4;      // single bit (hi==lo)
    u8,   unassigned, set_unassigned : 7, 5; // 3 bits
}

/// Represents an entry in the BGP Graceful Restart capability.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GracefulRestartEntry {
    afi: AddressFamilyId,
    safi: SubsequentAfi,
    /// Forwarding State (F) bit, which can be used to indicate whether the
    /// forwarding state for routes that were advertised with the given AFI
    /// and SAFI has indeed been preserved during the previous BGP restart.
    preserve_forwarding: bool,
}

#[cfg(test)]
mod tests {
    use std::net::Ipv4Addr;

    use bytes::BytesMut;
    use eyre::Result;

    use crate::{
        constants::{AddressFamilyId, SubsequentAfi},
        open::{Capability, ExtendedNextHop, GracefulRestartEntry, OpenMessage, OpenOption},
        parser::ParserContext,
    };

    macro_rules! test_capability_roundtrip {
        ($name:ident, $input_bytes: expr, $expected: expr) => {
            #[test]
            fn $name() -> Result<()> {
                let ctx = &default_v6_context();
                let (buf, parsed) = Capability::from_wire(ctx, $input_bytes)?;
                assert_eq!(parsed, $expected);
                assert!(buf.is_empty());
                let mut out = BytesMut::with_capacity(u16::MAX as usize);
                parsed.to_wire(ctx, &mut out)?;
                assert_eq!(out.to_vec(), $input_bytes);
                Ok(())
            }
        };
    }
    /// Creates a default context for evaluating the test cases below.
    /// This uses `four_octet_asn` set to true and address_family to IPv6.
    fn default_v6_context() -> ParserContext {
        ParserContext {
            four_octet_asn: Some(true),
            address_family: Some(AddressFamilyId::Ipv6),
        }
    }

    test_capability_roundtrip!(
        test_multiptotocol,
        &[0x01, 0x04, 0x00, 0x01, 0x00, 0x01],
        Capability::MultiProtocol {
            afi: AddressFamilyId::Ipv4,
            safi: SubsequentAfi::Unicast
        }
    );

    test_capability_roundtrip!(test_route_refresh, &[0x02, 0x00], Capability::RouteRefresh);

    test_capability_roundtrip!(
        test_route_refresh_cisco_unsupported,
        &[0x80, 0x00],
        Capability::UnknownCapability {
            type_code: 0x80,
            length: 0x00,
            value: vec![]
        }
    );

    test_capability_roundtrip!(
        test_four_octet_asn,
        &[0x41, 0x04, 0x00, 0x00, 0xfd, 0xe8],
        Capability::FourByteAsn { asn: 65000 }
    );

    test_capability_roundtrip!(
        test_graceful_restart,
        &[0x40, 0x06, 0x01, 0x2c, 0x00, 0x01, 0x01, 0x00],
        Capability::GracefulRestart {
            restart_state: false,
            restart_time: 300,
            entries: vec![GracefulRestartEntry {
                afi: AddressFamilyId::Ipv4,
                safi: SubsequentAfi::Unicast,
                preserve_forwarding: false,
            }]
        }
    );

    test_capability_roundtrip!(
        test_extended_nexthop,
        &[
            0x05, 0x12, 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00, 0x02, 0x00, 0x02,
            0x00, 0x01, 0x00, 0x80, 0x00, 0x02
        ],
        Capability::ExtendedNextHop {
            entries: vec![
                ExtendedNextHop {
                    afi: AddressFamilyId::Ipv4,
                    safi: SubsequentAfi::Unicast,
                    nh_afi: AddressFamilyId::Ipv6,
                },
                ExtendedNextHop {
                    afi: AddressFamilyId::Ipv4,
                    safi: SubsequentAfi::Multicast,
                    nh_afi: AddressFamilyId::Ipv6,
                },
                ExtendedNextHop {
                    afi: AddressFamilyId::Ipv4,
                    safi: SubsequentAfi::MplsLabelledVpn,
                    nh_afi: AddressFamilyId::Ipv6
                }
            ]
        }
    );

    #[test]
    fn test_open_roundtrip() -> Result<()> {
        let buf = &[
            0x04, 0xfd, 0xe8, 0x00, 0xb4, 0x0a, 0x00, 0x00, 0x01, 0x38, 0x02, 0x06, 0x01, 0x04,
            0x00, 0x01, 0x00, 0x01, 0x02, 0x02, 0x80, 0x00, 0x02, 0x02, 0x02, 0x00, 0x02, 0x06,
            0x41, 0x04, 0x00, 0x00, 0xfd, 0xe8, 0x02, 0x08, 0x40, 0x06, 0x01, 0x2c, 0x00, 0x01,
            0x01, 0x00, 0x02, 0x14, 0x05, 0x12, 0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01,
            0x00, 0x02, 0x00, 0x02, 0x00, 0x01, 0x00, 0x80, 0x00, 0x02,
        ];
        let expected = OpenMessage {
            version: 4,
            asn: 65000,
            hold_time: 180,
            identifier: Ipv4Addr::new(10, 0, 0, 1),
            options: vec![
                OpenOption::Capabilities(vec![Capability::MultiProtocol {
                    afi: AddressFamilyId::Ipv4,
                    safi: SubsequentAfi::Unicast,
                }]),
                OpenOption::Capabilities(vec![Capability::UnknownCapability {
                    type_code: 0x80,
                    length: 0x00,
                    value: vec![],
                }]),
                OpenOption::Capabilities(vec![Capability::RouteRefresh]),
                OpenOption::Capabilities(vec![Capability::FourByteAsn { asn: 65000 }]),
                OpenOption::Capabilities(vec![Capability::GracefulRestart {
                    restart_state: false,
                    restart_time: 300,
                    entries: vec![GracefulRestartEntry {
                        afi: AddressFamilyId::Ipv4,
                        safi: SubsequentAfi::Unicast,
                        preserve_forwarding: false,
                    }],
                }]),
                OpenOption::Capabilities(vec![Capability::ExtendedNextHop {
                    entries: vec![
                        ExtendedNextHop {
                            afi: AddressFamilyId::Ipv4,
                            safi: SubsequentAfi::Unicast,
                            nh_afi: AddressFamilyId::Ipv6,
                        },
                        ExtendedNextHop {
                            afi: AddressFamilyId::Ipv4,
                            safi: SubsequentAfi::Multicast,
                            nh_afi: AddressFamilyId::Ipv6,
                        },
                        ExtendedNextHop {
                            afi: AddressFamilyId::Ipv4,
                            safi: SubsequentAfi::MplsLabelledVpn,
                            nh_afi: AddressFamilyId::Ipv6,
                        },
                    ],
                }]),
            ],
        };

        let ctx = default_v6_context();
        let (_, parsed) = OpenMessage::from_wire(&ctx, buf)?;
        assert_eq!(parsed, expected);
        let mut out = BytesMut::with_capacity(u16::MAX as usize);
        parsed.to_wire(&ctx, &mut out)?;
        assert_eq!(buf.to_vec(), out.to_vec());
        Ok(())
    }
}