1 files changed, 47 insertions, 0 deletions
diff --git a/src/gens.rs b/src/gens.rs
new file mode 100644
index 000000000000..cd0af16754f3
--- /dev/null
+++ b/src/gens.rs
@@ -0,0 +1,47 @@
+//! Helpers to determine if a node can send a particular packet.
+
+/// These functions are unfortunately all a bit stringly typed atm
+pub mod can_send {
+    use crate::data::{Level, PacketType, Upgrade};
+    use Level::*;
+
+    #[inline]
+    pub fn base() -> Vec<&'static str> {
+        vec!["ping", "capture"]
+    }
+
+    /// Guard nodes are the most versatile
+    #[inline]
+    pub fn guard(lvl: Level) -> Vec<&'static str> {
+        match lvl {
+            // This is just kinda gross
+            One => {
+                let mut p = base();
+                p.append(&mut vec!["payload", "cns", "reset"]);
+                p
+            }
+            Two => {
+                let mut p = guard(One);
+                p.append(&mut vec!["nitm", "virus"]);
+                p
+            }
+            Three => {
+                let mut p = guard(Two);
+                p.append(&mut vec!["takeover"]);
+                p
+            }
+        }
+    }
+
+    /// Compute nodes can sennd compute packets
+    pub fn compute() -> Vec<&'static str> {
+        let mut p = base();
+        p.append(&mut vec!["compute"]);
+        p
+    }
+
+    /// Relays only relay!
+    pub fn relay() -> Vec<&'static str> {
+        vec![]
+    }
+}