(* SPDX-License-Identifier: LGPL-3.0-only *)

module U64 = Stdint.Uint64
module U32 = Stdint.Uint32

type sid =
  { sid_ident_auth : U64.t (* 6 B *)
  ; sid_sub_auths  : U32.t array (* max. 15 × *)
  }
and sub_auths = U32.t array

let sizeof_ident_auth = 6
let max_ident_auth    = U64.of_string "0x0000_ffff_ffff_ffff"
let sizeof_sub_auth   = 4
let max_subauth_count = 15

let create_unsafe sa ia =
  { sid_ident_auth = ia
  ; sid_sub_auths  = sa }

(* There isn’t much to validate to begin with except for the hard cap on
   the number of subauths. *)
let create ?(sa=[||]) ia =
  if Array.length sa > max_subauth_count then None else
  Some (create_unsafe sa ia)

let get_ident_auth s = s.sid_ident_auth
let get_sub_auths s = s.sid_sub_auths

exception Nope

let equal_sub_auths saa sab =
  try
    Array.iter2
      (fun saa sab -> if U32.compare saa sab <> 0 then raise Nope)
      saa sab;
    true
  with Nope -> false

let equal a b =
  U64.compare a.sid_ident_auth b.sid_ident_auth = 0
  && Array.length a.sid_sub_auths = Array.length b.sid_sub_auths
  && equal_sub_auths a.sid_sub_auths b.sid_sub_auths

module StringFmt = struct

  let expect_char s c p =
    if s.[p] = c then () else
    raise
      (Invalid_argument
         (Printf.sprintf
            "Invalid SID [%s]: expected ‘%c’ at position %d, found ‘%c’"
            s c p s.[p]))

  let is_digit c = '0' <= c && c <= '9'

  let read_decimal_string f s p =
    let n = String.length s in
    assert (p < n);
    let p' = ref p in
    let b = Buffer.create 16 in
    while !p' < n && is_digit s.[!p'] do
      Buffer.add_char b s.[!p'];
      incr p'
    done;
    let nb = Buffer.length b in
    if nb = 0 then
      raise (Invalid_argument
               (Printf.sprintf
                  "Invalid SID [%s]: expected decimal at position %d" s p))
    else
      p + nb,
      f (Buffer.contents b)

  let read_decimal_u64 = read_decimal_string U64.of_string
  let read_decimal_u32 = read_decimal_string U32.of_string

  (*
   * The spec ([MS-DTYP]):
   *
   *    2.4.2.1: SID= "S-1-" IdentifierAuthority 1*SubAuthority
   *)
  let decode s =
    let n = String.length s in
    if n <= 4 then
      Error
        (Printf.sprintf
           "Invalid SID: ‘%s’ too short to be a SID in string format" s)
    else begin
      try
        expect_char s 'S' 0;
        expect_char s '-' 1;
        expect_char s '1' 2;
        expect_char s '-' 3;
        let p = 4 in
        let p, ia = read_decimal_u64 s p in
        if ia > max_ident_auth then
          raise (Invalid_argument
                   (Printf.sprintf
                      "Invalid SID: identifier authority cannot fit 6 B (%s)"
                      (U64.to_string max_ident_auth)));
        let sa = ref [] and p' = ref p in
        while !p' < n - 1 && List.length !sa < max_subauth_count do
          expect_char s '-' !p';
          let np, d =
            try read_decimal_u32 s (!p' + 1)
            with Invalid_argument e ->
              (* Brr, but Stdint’s error messages aren’t overly instructive. *)
              raise (Invalid_argument
                       (Printf.sprintf
                          "Invalid SID: error parsing subauth at position %d, \
                           (err: %s)"
                          (!p' + 1) e))
          in
          sa := d :: !sa;
          p' := np
        done;
        Ok { sid_ident_auth = ia
           ; sid_sub_auths  = Array.of_list (List.rev !sa)
           }
      with Invalid_argument e -> Error e;
    end

  let fmt_ident_auth b ia =
    Buffer.add_string b (U64.to_string ia)

  let fmt_sub_auths b sas =
    Array.iter
      (fun sa ->
         Buffer.add_char b '-';
         Buffer.add_string b (U32.to_string sa))
      sas

  let encode s =
    let b = Buffer.create 16 in
    Buffer.add_string b "S-1-";
    fmt_ident_auth b s.sid_ident_auth;
    fmt_sub_auths  b s.sid_sub_auths;
    Buffer.contents b

end (* [module StringFmt] *)

module PacketRep = struct (* [MS-DTYP] 2.4.22 *)

  type endian = Big | Little

  let wordlen        = 4 (* sizeof int *)
  let min_pktrep_len = 1 + 1 + sizeof_ident_auth
  let max_pktrep_len = 1 + 1
                     + sizeof_ident_auth
                     + max_subauth_count * sizeof_sub_auth
  let pktrep_sa_off  = min_pktrep_len

  let encode ?(endian=Little) s =
    let nsa = Array.length s.sid_sub_auths in
    let l = 8 + nsa * sizeof_sub_auth in
    let b = Bytes.create l in
    let o = ref 0 in
    let pushbyte c = char_of_int c |> Bytes.set b !o; incr o in

    assert (0 <= nsa && nsa <= 15);

    pushbyte 1;
    pushbyte nsa;

    let getia n =
      pushbyte (U64.to_int (U64.shift_right s.sid_ident_auth n) land 0xff)
    in (* big endian!, cf. [MS-DTYP] 2.4.1.1 *)
    getia 5; getia 4; getia 3; getia 2; getia 1; getia 0;

    let write_u32 =
      match endian with
        | Big    -> U32.to_bytes_big_endian
        | Little -> U32.to_bytes_little_endian
    in
    Array.iteri
      (fun i sa ->
         let o' = !o + i * wordlen in
         write_u32 sa b o')
      s.sid_sub_auths;
    b

  let decode ?(endian=Little) b =
    let l = Bytes.length b in
    if l < min_pktrep_len || max_pktrep_len < l then
      Error (Printf.sprintf
               "bad input size: expected %d–%d B, got %d B"
               min_pktrep_len max_pktrep_len l) else
    if l mod wordlen <> 0 then
      Error (Printf.sprintf
               "bad input size: not divisible by word length (%d)"
               wordlen) else
    let v = Bytes.get b 0 |> int_of_char in
    if v <> 0x01 then
      Error (Printf.sprintf
               "input malformed: expected SID version=0x01, got 0x%0.2x" v) else
    let nsa = Bytes.get b 1 |> int_of_char in
    if max_subauth_count < nsa then
      Error (Printf.sprintf
               "input malformed: up to %d subAuthority elements permitted, \
                %d specified"
               max_subauth_count nsa) else
    let getbyte n ia = (* b[n] << (5 - (n - 2)) *)
      U64.logor ia
          (U64.shift_left (Bytes.get b n |> int_of_char |> U64.of_int) (5 - (n - 2)))
    in
    let ia = U64.zero
             |> getbyte 2 |> getbyte 3 |> getbyte 4
             |> getbyte 5 |> getbyte 6 |> getbyte 7 in
    let sas = Array.make nsa (U32.zero)
    and read_u32 =
      match endian with
        | Big    -> U32.of_bytes_big_endian
        | Little -> U32.of_bytes_little_endian
    in
    for i = 0 to (nsa - 1) do
      let off = pktrep_sa_off + i * sizeof_sub_auth in
      sas.(i) <- read_u32 b off
    done;
    Ok { sid_ident_auth = ia
       ; sid_sub_auths  = sas
       }

end (* [module PacketRep] *)

module WellKnown = struct
  (*
   * see also
   * https://docs.microsoft.com/en-us/windows/desktop/secauthz/well-known-sids
   *)

  let null = create_unsafe [| U32.zero |] U64.zero
  let everyone = create_unsafe [| U32.zero |] U64.one
  let world = everyone
  let local = create_unsafe [| U32.zero |] (U64.of_int 2)
  let creator_owner_id = create_unsafe [| U32.zero |] (U64.of_int 3)
  let creator_group_id = create_unsafe [| U32.one |] (U64.of_int 3)
  let elite = create_unsafe [| U32.of_int 3 ; U32.of_int 3; U32.of_int 7 |] U64.one

  module Prefix = struct
    type toplevel_auth = ?sa:sub_auths -> unit -> sid

    let security_null_sid_authority    ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x00)
    let security_world_sid_authority   ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x01)
    let security_local_sid_authority   ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x02)
    let security_creator_sid_authority ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x03)
    let security_nt_authority          ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x05)
    let security_app_package_authority      ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x0f)
    let security_mandatory_label_authority  ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x10)
    let security_scoped_policy_id_authority ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x11)
    let security_authentication_authority   ?(sa=[||]) () = create_unsafe sa (U64.of_int 0x12)
  end
end

let of_string = StringFmt.decode
let to_string = StringFmt.encode

type t = sid