use std::{
collections::{BTreeMap, BTreeSet},
net::Ipv4Addr,
time::Duration,
};
use crate::{
bgp::BgpRoute,
network::Network,
ospf::OspfImpl,
types::{Ipv4Prefix, Prefix, PrefixMap, RouterId, ASN},
};
use super::{Addressor, CfgGen, ExportError};
use ipnet::Ipv4Net;
use itertools::Itertools;
use maplit::btreemap;
pub const RUNNER_PREAMBLE: &str =
"#!/usr/bin/env python3\n\nimport sys\nimport time\n\n\ntime.sleep(5)\n\n";
pub const RUNNER_POSTAMBLE: &str = "\nwhile True:\n time.sleep(1)\n";
#[derive(Debug)]
pub struct ExaBgpCfgGen<P: Prefix> {
router: RouterId,
asn: ASN,
routes: BTreeMap<P, BTreeMap<Duration, Option<BgpRoute<P>>>>,
neighbors: BTreeSet<RouterId>,
current_time: Duration,
}
impl<P: Prefix> ExaBgpCfgGen<P> {
pub fn new<Q, Ospf: OspfImpl>(
net: &Network<P, Q, Ospf>,
router: RouterId,
) -> Result<Self, ExportError> {
let r = net.get_device(router)?.external_or_err()?;
Ok(Self {
router,
asn: r.asn(),
routes: r
.active_routes
.iter()
.map(|(p, r)| (*p, btreemap! {Duration::ZERO => Some(r.clone())}))
.collect(),
neighbors: r.neighbors.iter().copied().collect(),
current_time: Duration::ZERO,
})
}
pub fn step_time(&mut self, step: Duration) {
self.current_time += step;
}
pub fn generate_script<A: Addressor<P>>(
&self,
addressor: &mut A,
) -> Result<String, ExportError> {
let script = String::from(RUNNER_PREAMBLE);
Ok(script + &self.generate_script_no_loop(addressor)?)
}
pub fn generate_lines<A: Addressor<P>>(
&self,
addressor: &mut A,
) -> Result<Vec<(Vec<String>, Duration)>, ExportError> {
let neighbors = self
.neighbors
.iter()
.map(|x| addressor.iface_address(*x, self.router))
.collect::<Result<Vec<Ipv4Addr>, ExportError>>()?
.into_iter()
.map(|x| format!("neighbor {x}"))
.join(", ");
let mut result = Vec::new();
let mut times_routes: BTreeMap<_, Vec<_>> = Default::default();
for (p, routes) in self.routes.iter() {
for (time, route) in routes.iter() {
times_routes
.entry(*time)
.or_default()
.push((*p, route.as_ref()));
}
}
for (time, routes) in times_routes {
let mut ads: Vec<String> = Vec::new();
for (p, r) in routes {
for net in addressor.prefix(p)? {
if let Some(r) = r {
let r = r.clone().with_prefix(Ipv4Prefix::from(net));
ads.push(format!(
"sys.stdout.write(\"{neighbors} {}\\n\")",
announce_route(&r)
))
} else {
ads.push(format!(
"sys.stdout.write(\"{neighbors} {}\\n\")",
withdraw_route(net)
))
}
}
}
result.push((ads, time));
}
Ok(result)
}
fn generate_script_no_loop<A: Addressor<P>>(
&self,
addressor: &mut A,
) -> Result<String, ExportError> {
let lines = self.generate_lines(addressor)?;
let mut script = String::new();
let mut current_time = Duration::ZERO;
for (routes, time) in lines {
if !time.is_zero() {
script.push_str(&format!(
"time.sleep({})\n",
(time - current_time).as_secs_f64()
));
}
current_time = time;
for route in routes {
script.push_str(&route);
script.push('\n');
}
script.push_str("sys.stdout.flush()\n");
}
script.push_str(RUNNER_POSTAMBLE);
Ok(script)
}
fn generate_neighbor_cfg<A: Addressor<P>, Q, Ospf: OspfImpl>(
&self,
net: &Network<P, Q, Ospf>,
addressor: &mut A,
neighbor: RouterId,
) -> Result<String, ExportError> {
let asn = net.get_device(neighbor)?.asn();
Ok(format!(
"\
neighbor {} {{
router-id {};
local-address {};
local-as {};
peer-as {};
family {{ ipv4 unicast; }}
capability {{ route-refresh; }}
}}",
addressor.iface_address(neighbor, self.router)?,
addressor.router_address(self.router)?,
addressor.iface_address(self.router, neighbor)?,
self.asn.0,
asn.0,
))
}
pub fn neighbors(&self) -> &BTreeSet<RouterId> {
&self.neighbors
}
}
pub fn announce_route<P: Prefix>(route: &BgpRoute<P>) -> String {
let prefix: Ipv4Net = route.prefix.into();
format!(
"announce route {prefix} next-hop self as-path [{}]{}{}",
route.as_path.iter().map(|x| x.0).join(", "),
if let Some(med) = route.med {
format!(" metric {med}")
} else {
String::new()
},
if route.community.is_empty() {
String::new()
} else {
format!(
" extended-community [{}]",
route
.community
.iter()
.map(|x| format!("{}:{}", 65535, x)) .join(", ")
)
},
)
}
pub fn withdraw_route(prefix: Ipv4Net) -> String {
format!("withdraw route {prefix}")
}
impl<P: Prefix, A: Addressor<P>, Q, Ospf: OspfImpl> ExternalCfgGen<P, Q, Ospf, A>
for ExaBgpCfgGen<P>
{
fn generate_config(
&mut self,
net: &Network<P, Q, Ospf>,
addressor: &mut A,
) -> Result<String, ExportError> {
Ok(self
.neighbors
.iter()
.map(|x| self.generate_neighbor_cfg(net, addressor, *x))
.collect::<Result<Vec<String>, ExportError>>()?
.into_iter()
.join("\n"))
}
fn advertise_route(
&mut self,
_net: &Network<P, Q, Ospf>,
addressor: &mut A,
route: &BgpRoute<P>,
) -> Result<String, ExportError> {
self.routes
.entry(route.prefix)
.or_default()
.insert(self.current_time, Some(route.clone()));
self.generate_script(addressor)
}
fn withdraw_route(
&mut self,
_net: &Network<P, Q, Ospf>,
addressor: &mut A,
prefix: P,
) -> Result<String, ExportError> {
self.routes
.entry(prefix)
.or_default()
.insert(self.current_time, None);
self.generate_script(addressor)
}
fn establish_ebgp_session(
&mut self,
net: &Network<P, Q, Ospf>,
addressor: &mut A,
neighbor: RouterId,
) -> Result<String, ExportError> {
self.neighbors.insert(neighbor);
self.generate_config(net, addressor)
}
fn teardown_ebgp_session(
&mut self,
net: &Network<P, Q, Ospf>,
addressor: &mut A,
neighbor: RouterId,
) -> Result<String, ExportError> {
self.neighbors.remove(&neighbor);
self.generate_config(net, addressor)
}
}