let () = Mirage_crypto_rng_unix.initialize ()

let csr_privkey =
  lazy
    (`RSA (Mirage_crypto_pk.Rsa.generate ~bits:1024 ()))

let ca_privkey =
  lazy
    (`RSA (Mirage_crypto_pk.Rsa.generate ~bits:1024 ()))

let organization = "Free Shell Servers, Inc."

let dn_of_name ?(organization=organization) name =
  let open X509.Distinguished_name in
  let open Relative_distinguished_name in
  [ singleton (CN name) |> add (O organization) ]

let ca_cert =
  let gen_ca () =
    let valid_from = Ptime_clock.now () in
    let valid_until =
      Ptime.add_span valid_from
        (Ptime.Span.of_int_s 60)
      |> Option.get
    in
    let subject = dn_of_name "Bad Authority, Inc." in
    let extensions =
      let open X509.Extension in
      empty
      |> add Key_usage
        (true, [`Key_cert_sign; `CRL_sign; `Digital_signature; `Content_commitment])
      |> add Basic_constraints (true, (true, None))
    in
    let csr =
      X509.Signing_request.create
        subject
        (Lazy.force ca_privkey)
    in
    let csr =
      Result.fold csr ~ok:Fun.id
        ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
    in
    X509.Signing_request.sign
      csr
      ~valid_from
      ~valid_until
      ~extensions
      (Lazy.force ca_privkey)
      subject
  in
  lazy (Result.get_ok (gen_ca ()))

let csr subject =
  X509.Signing_request.create
    subject
    (Lazy.force csr_privkey)

let check_csr_dn_good () =
  let subject = dn_of_name "reynir" in
  let csr = csr subject in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  Alcotest.(check @@ result unit Alcotest.reject)
    "good dn in csr"
    (Ok ())
    (Cert_service.check_csr_dn
       csr subject)

let check_csr_dn_different () =
  let subject = dn_of_name "reynir" in
  let csr = csr subject in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  match Cert_service.check_csr_dn csr (dn_of_name "notreynir") with
  | Ok () -> Alcotest.fail "check succeeded, expected failure"
  | Error (`Msg _) -> ()

let check_csr_dn_diff_org () =
  let subject = dn_of_name ~organization:"Evil Corp" "reynir" in
  let csr = csr subject in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  match Cert_service.check_csr_dn csr (dn_of_name "reynir") with
  | Ok () -> Alcotest.fail "check succeeded, expected failure"
  | Error (`Msg _) -> ()

let check_csr_dn_extra () =
  let subject = dn_of_name "reynir" in
  let csr = csr (subject @ dn_of_name "bob") in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  match Cert_service.check_csr_dn csr (dn_of_name "reynir") with
  | Ok () -> Alcotest.fail "check succeeded, expected failure"
  | Error (`Msg _) -> ()

let check_csr_tests = [
  "check_csr_dn_good", `Quick, check_csr_dn_good;
  "check_csr_dn_different", `Quick, check_csr_dn_different;
  "check_csr_dn_diff_org", `Quick, check_csr_dn_diff_org;
  "check_csr_dn_extra", `Quick, check_csr_dn_extra;
]

let good_sign () =
  let t = {
    Cert_service.host = "example.com";
    organization;
    cacert = Lazy.force ca_cert;
    cakey = Lazy.force ca_privkey;
  }
  in
  let name = "reynir" in
  let csr = csr (dn_of_name name) in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  match Cert_service.sign t csr name with
  | Ok _cert -> ()
  | Error (`Msg e) -> Alcotest.failf "cert sign failed unexpectedly: %s" e

let bad_sign () =
  let t = {
    Cert_service.host = "example.com";
    organization;
    cacert = Lazy.force ca_cert;
    cakey = Lazy.force ca_privkey;
  }
  in
  let name = "reynir" in
  let csr = csr (dn_of_name name) in
  let csr =
    Result.fold csr ~ok:Fun.id
      ~error:(fun (`Msg e) -> Alcotest.failf "CSR error %s" e)
  in
  match Cert_service.sign t csr "notreynir" with
  | Ok cert ->
    Alcotest.failf
      "cert sign succeeded unexpectedly:\n%a" X509.Certificate.pp
      cert
  | Error (`Msg _) -> ()

let cert_tests = [
  "good_sign", `Slow, good_sign;
  "bad_sign", `Slow, bad_sign;
]

let tests : unit Alcotest.test list= [
  "check_csr_dn", check_csr_tests;
  "cert_tests", cert_tests;
]

let () =
  Alcotest.run
    "Cert-service tests"
    tests