(*
 * Copyright (c) 2012 Anil Madhavapeddy <anil@recoil.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *)

(* Convert a UNIX /etc/services into an ML module to lookup entries *)
open Printf

let hashtbl_add_list h k v =
  try
    let l = Hashtbl.find h k in
    l := v :: !l
  with Not_found -> Hashtbl.add h k (ref [v])

let spaced_list = Stringext.split_trim_left ~on:" " ~trim:" "
let nonempty = List.filter ((<>) "")

let iter f h =
  let bindings = Hashtbl.fold (fun k v a -> (k,v)::a) h [] in
  List.iter (fun (k, v) -> f k v) (List.sort compare bindings)

let _ =
  let fin = open_in Sys.argv.(1) in
  let tcp_ports = Hashtbl.create 1 in
  let udp_ports = Hashtbl.create 1 in
  let ports_tcp = Hashtbl.create 1 in
  let ports_udp = Hashtbl.create 1 in
  let tcp_services = Hashtbl.create 1 in
  let udp_services = Hashtbl.create 1 in
  (try while true do
    let line = input_line fin in
    match line.[0] with
    |'#'|' ' -> ()
    |_ ->
      Scanf.sscanf line "%s %d/%s %s@\n" (fun svc port proto rest ->
        let alias_s = List.hd (Stringext.split ~on:'#' (" "^rest)) in
        let aliases = nonempty (spaced_list alias_s) in
        match proto with
        |"tcp" ->
          List.iter (fun svc ->
            hashtbl_add_list tcp_ports svc port;
            hashtbl_add_list ports_tcp port svc;
            Hashtbl.replace tcp_services svc ()
          ) (svc::aliases)
        |"udp" ->
          List.iter (fun svc ->
            hashtbl_add_list udp_ports svc port;
            hashtbl_add_list ports_udp port svc;
            Hashtbl.replace udp_services svc ();
          ) (svc::aliases)
        |"ddp" | "sctp" | "divert" -> ()
        |x -> failwith ("unknown proto " ^ x)
      )
  done with End_of_file -> ());
  let print_keys quote ppf table =
    iter (fun k _v -> fprintf ppf ("%s; ") (quote k)) table in
  let print_values quote ppf table =
    iter (fun _k v -> fprintf ppf "[ %s ]; "
             (String.concat "; " (List.map quote !v))) table in
  let quote_string s = sprintf "%S" s in
  printf "(* Autogenerated by gen_services.ml, do not edit directly *)\n";

  printf "let tcp_port_of_service_tables = (\n  [| %a |],\n  [| %a |]\n)\n\n"
    (print_keys quote_string) tcp_ports
    (print_values string_of_int) tcp_ports;

  printf "let udp_port_of_service_tables = (\n  [| %a |],\n  [| %a |]\n)\n\n"
    (print_keys quote_string) udp_ports
    (print_values string_of_int) udp_ports;

  printf "let service_of_tcp_port_tables = (\n  [| %a |],\n  [| %a |]\n)\n\n"
    (print_keys string_of_int) ports_tcp
    (print_values quote_string) ports_tcp;

  printf "let service_of_udp_port_tables = (\n  [| %a |],\n  [| %a |]\n)\n\n"
    (print_keys string_of_int) ports_udp
    (print_values quote_string) ports_udp;

  let hashset_elems table =
    Hashtbl.fold (fun k () a -> quote_string k :: a) table []
    |> List.sort String.compare
    |> String.concat "; "
  in
  printf "let known_tcp_services =\n  [ %s ]\n\n"
    (hashset_elems tcp_services);
  printf "let known_udp_services =\n [ %s ]\n\n"
    (hashset_elems udp_services);

  printf "let known_services = [\n";
  printf "  (\"tcp\", known_tcp_services);\n";
  printf "  (\"udp\", known_udp_services) ]\n\n";