open Utils open Ipaddr let assoc_one key value = `Assoc [ (key, value) ] module Payload = struct module Udp = struct type _ t = Sport : int t | Dport : int t let to_string : type a. a t -> string = function | Sport -> "sport" | Dport -> "dport" end module Tcp = struct type _ t = Sport : int t | Dport : int t let to_string : type a. a t -> string = function | Sport -> "sport" | Dport -> "dport" end module Ipv4 = struct type _ t = Saddr : V4.Prefix.t t | Daddr : V4.Prefix.t t let to_string : type a. a t -> string = function | Saddr -> "saddr" | Daddr -> "daddr" end module Ipv6 = struct type _ t = Saddr : V6.Prefix.t t | Daddr : V6.Prefix.t t let to_string : type a. a t -> string = function | Saddr -> "saddr" | Daddr -> "daddr" end type _ t = | Udp : 'a Udp.t -> 'a t | Tcp : 'a Tcp.t -> 'a t | Ipv4 : 'a Ipv4.t -> 'a t | Ipv6 : 'a Ipv6.t -> 'a t let to_json (type a) (payload : a t) = let protocol, field = match payload with | Udp udp -> ("udp", Udp.to_string udp) | Tcp tcp -> ("tcp", Tcp.to_string tcp) | Ipv4 ipv4 -> ("ip", Ipv4.to_string ipv4) | Ipv6 ipv6 -> ("ip6", Ipv6.to_string ipv6) in assoc_one "payload" (`Assoc [ ("protocol", `String protocol); ("field", `String field) ]) end module Expr = struct type _ t = | String : string -> string t | Number : int -> int t | Boolean : bool -> int t | Ipv4 : V4.Prefix.t -> V4.Prefix.t t | Ipv6 : V6.Prefix.t -> V6.Prefix.t t | List : 'a t list -> 'a t | Set : 'a t list -> 'a t | Range : 'a t * 'a t -> 'a t | Payload : 'a Payload.t -> 'a t let rec to_json : type a. a t -> Yojson.Basic.t = function | String str -> `String str | Number num -> `Int num | Boolean bool -> `Bool bool | Ipv4 ipv4 -> let network = V4.(to_string (Prefix.network ipv4)) in let len = V4.Prefix.bits ipv4 in assoc_one "prefix" (`Assoc [ ("addr", `String network); ("len", `Int len) ]) | Ipv6 ipv6 -> let network = V6.(to_string (Prefix.network ipv6)) in let len = V6.Prefix.bits ipv6 in assoc_one "prefix" (`Assoc [ ("addr", `String network); ("len", `Int len) ]) | List list -> `List (List.map to_json list) | Set set -> assoc_one "set" (`List (List.map to_json set)) | Range (a, b) -> assoc_one "range" (`List [ to_json a; to_json b ]) | Payload payload -> Payload.to_json payload end module Match = struct type (_, _) t = Eq : ('a, 'a) t | NotEq : ('a, 'a) t let to_string : type a b. (a, b) t -> string = function | Eq -> "==" | NotEq -> "!=" let to_json (type a b) (op : (a, b) t) = `String (to_string op) end module Counter = struct type t = | NamedCounter of string | AnonCounter of { packets : int; bytes : int } let to_json = function | NamedCounter n -> assoc_one "counter" (`String n) | AnonCounter { packets; bytes } -> assoc_one "counter" (`Assoc [ ("packets", `Int packets); ("bytes", `Int bytes) ]) end module Verdict = struct type t = Accept | Drop | Continue | Return | Jump of string | Goto of string let to_json = function | Accept -> assoc_one "accept" `Null | Drop -> assoc_one "drop" `Null | Continue -> assoc_one "continue" `Null | Return -> assoc_one "return" `Null | Jump s -> assoc_one "jump" (assoc_one "target" (`String s)) | Goto s -> assoc_one "goto" (assoc_one "target" (`String s)) end module Stmt = struct type _ t = | Match : (('a, 'b) Match.t * 'a Expr.t * 'b Expr.t) -> unit t | Counter : Counter.t -> unit t | Verdict : Verdict.t -> unit t | NoTrack : unit t | Log : { prefix : string option; group : int option } -> unit t let to_json : type a. a t -> Yojson.Basic.t = function | Match (op, left, right) -> assoc_one "match" (`Assoc [ ("left", Expr.to_json left); ("right", Expr.to_json right); ("op", Match.to_json op); ]) | Counter counter -> Counter.to_json counter | Verdict verdict -> Verdict.to_json verdict | NoTrack -> assoc_one "notrack" `Null | Log { prefix; group } -> let elems = [ Option.map (fun p -> ("prefix", `String p)) prefix; Option.map (fun g -> ("group", `Int g)) group; ] in assoc_one "log" (`Assoc (deoptionalise elems)) end module Family = struct type t = Ipv4 | Ipv6 | Inet let to_string = function Ipv4 -> "ip" | Ipv6 -> "ip6" | Inet -> "inet" let to_json f = `String (to_string f) end module Table = struct type t = { family : Family.t; name : string } let to_json { family; name } = assoc_one "table" (`Assoc [ ("family", Family.to_json family); ("name", `String name) ]) end module Chain = struct type t = { family : Family.t; table : string; name : string } let to_json { family; table; name } = assoc_one "chain" (`Assoc [ ("family", Family.to_json family); ("table", `String table); ("name", `String name); ]) end module Rule = struct type t = { family : Family.t; table : string; chain : string; expr : unit Stmt.t list; } let to_json { family; table; chain; expr } = assoc_one "rule" (`Assoc [ ("family", Family.to_json family); ("table", `String table); ("chain", `String chain); ("expr", `List (List.map Stmt.to_json expr)); ]) end module Command = struct type t = | FlushRuleset | AddTable of Table.t | FlushTable of Table.t | AddChain of Chain.t | FlushChain of Chain.t | AddRule of Rule.t let to_json = function | FlushRuleset -> assoc_one "flush" (assoc_one "ruleset" `Null) | AddTable table -> assoc_one "add" (Table.to_json table) | FlushTable table -> assoc_one "flush" (Table.to_json table) | AddChain chain -> assoc_one "add" (Chain.to_json chain) | FlushChain chain -> assoc_one "flush" (Chain.to_json chain) | AddRule rule -> assoc_one "add" (Rule.to_json rule) end let to_json commands = assoc_one "nftables" (`List (List.map Command.to_json commands))