/*
  webid-oidc, implementation of the Solid specification
  Copyright (C) 2020, 2021  Vivien Kraus

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as
  published by the Free Software Foundation, either version 3 of the
  License, or (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#include <utilities.h>

#define _(s) dgettext (PACKAGE, s)

SCM_SYMBOL (rs256, "RS256");
SCM_SYMBOL (rs512, "RS512");
SCM_SYMBOL (es256, "ES256");
SCM_SYMBOL (es384, "ES384");
SCM_SYMBOL (es512, "ES512");
SCM_SYMBOL (ps256, "PS256");
SCM_SYMBOL (ps384, "PS384");
SCM_SYMBOL (ps512, "PS512");

SCM_SYMBOL (sha256, "SHA-256");
SCM_SYMBOL (sha384, "SHA-384");
SCM_SYMBOL (sha512, "SHA-512");

SCM_SYMBOL (kty_ec, "EC");
SCM_SYMBOL (kty_rsa, "RSA");

SCM_SYMBOL (incompatible_alg, "incompatible-alg");

/* Required for utilities.h */
SCM_SYMBOL (p256, "P-256");
SCM_SYMBOL (p384, "P-384");
SCM_SYMBOL (p521, "P-521");
SCM_SYMBOL (kcrv, "crv");
SCM_SYMBOL (kx, "x");
SCM_SYMBOL (ky, "y");
SCM_SYMBOL (kd, "d");
SCM_SYMBOL (kn, "n");
SCM_SYMBOL (ke, "e");
SCM_SYMBOL (kp, "p");
SCM_SYMBOL (kq, "q");
SCM_SYMBOL (kdp, "dp");
SCM_SYMBOL (kdq, "dq");
SCM_SYMBOL (kqi, "qi");

void webid_oidc_random (size_t len, uint8_t * dst);
SCM webidoidc_hash_g (SCM alg, SCM payload);
SCM webidoidc_kty_g (SCM key);

static void
generate_random (void *unused, size_t len, uint8_t * dst)
{
  (void) unused;
  webid_oidc_random (len, dst);
}

static void
do_dsa_signature_clear (void *ptr)
{
  struct dsa_signature *sig = ptr;
  dsa_signature_clear (sig);
}

static void
dynwind_dsa_signature_clear (struct dsa_signature *sig)
{
  scm_dynwind_unwind_handler (do_dsa_signature_clear, sig,
			      SCM_F_WIND_EXPLICITLY);
}

SCM_DEFINE (webidoidc_jws_sign_g, "sign", 3, 0, 0,
	    (SCM alg, SCM key, SCM payload),
	    "Sign @var{payload} with @var{key}, using the given algorithm @var{alg}. If @var{alg} is not @code{'RS256}, @code{'RS512}, @code{'ES256}, @code{'ES384}, @code{'ES512}, @code{'PS256}, @code{'PS384}, or @code{'PS512}, or is not compatible with the key type, throw an error.")
{
  SCM digest;
  size_t c_digest_size;
  uint8_t *c_digest;
  if (scm_is_eq (alg, rs256)
      || scm_is_eq (alg, es256) || scm_is_eq (alg, ps256))
    {
      digest = webidoidc_hash_g (sha256, payload);
    }
  else if (scm_is_eq (alg, es384) || scm_is_eq (alg, ps384))
    {
      digest = webidoidc_hash_g (sha384, payload);
    }
  else if (scm_is_eq (alg, rs512)
	   || scm_is_eq (alg, es512) || scm_is_eq (alg, ps512))
    {
      digest = webidoidc_hash_g (sha512, payload);
    }
  else
    {
      scm_throw (incompatible_alg, scm_list_2 (alg, key));
    }
  c_digest = get_as_bytevector (digest, &c_digest_size, 1);
  if (scm_is_eq (alg, es256)
      || scm_is_eq (alg, es384) || scm_is_eq (alg, es512))
    {
      struct ecc_scalar c_key;
      SCM crv = scm_assq_ref (key, kcrv);
      const struct ecc_curve *c_crv = do_ecc_curve_load (crv, 1);
      struct dsa_signature sig;
      SCM r, s;
      uint8_t *c_r, *c_s, *c_ret;
      size_t c_r_size, c_s_size;
      scm_dynwind_begin (0);
      ecc_scalar_init (&c_key, c_crv);
      dynwind_ecc_scalar_clear (&c_key);
      dsa_signature_init (&sig);
      dynwind_dsa_signature_clear (&sig);
      if (!do_ecc_scalar_load (&c_key, key))
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      ecdsa_sign (&c_key, NULL, &generate_random, c_digest_size, c_digest,
		  &sig);
      r = wrap_mpz_t (sig.r);
      s = wrap_mpz_t (sig.s);
      c_r = get_as_bytevector (r, &c_r_size, 1);
      c_s = get_as_bytevector (s, &c_s_size, 1);
      scm_dynwind_end ();
      c_ret = scm_gc_malloc_pointerless (c_r_size + c_s_size, "signature");
      memcpy (c_ret, c_r, c_r_size);
      memcpy (c_ret + c_r_size, c_s, c_s_size);
      return wrap_bytevector (c_r_size + c_s_size, c_ret);
    }
  else
    {
      struct rsa_public_key c_pub;
      struct rsa_private_key c_key;
      mpz_t c_sig;
      SCM ret;
      scm_dynwind_begin (0);
      rsa_public_key_init (&c_pub);
      dynwind_rsa_public_key_clear (&c_pub);
      rsa_private_key_init (&c_key);
      dynwind_rsa_private_key_clear (&c_key);
      if (!do_rsa_public_key_load (&c_pub, key)
	  || !do_rsa_private_key_load (&c_key, key))
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      mpz_init (c_sig);
      dynwind_mpz_t_clear (&c_sig);
      if (scm_is_eq (alg, rs256)
	  && rsa_sha256_sign_digest_tr (&c_pub, &c_key, NULL,
					&generate_random, c_digest, c_sig))
	{
	  ret = wrap_mpz_t (c_sig);
	}
      else if (scm_is_eq (alg, rs512)
	       && rsa_sha512_sign_digest_tr (&c_pub, &c_key, NULL,
					     &generate_random, c_digest,
					     c_sig))
	{
	  ret = wrap_mpz_t (c_sig);
	}
      else if (scm_is_eq (alg, ps256))
	{
	  uint8_t padding[SHA256_DIGEST_SIZE];
	  webid_oidc_random (SHA256_DIGEST_SIZE, padding);
	  if (rsa_pss_sha256_sign_digest_tr
	      (&c_pub, &c_key, NULL, &generate_random, SHA256_DIGEST_SIZE,
	       padding, c_digest, c_sig))
	    {
	      ret = wrap_mpz_t (c_sig);
	    }
	  else
	    {
	      scm_wrong_type_arg ("sign", SCM_ARG2, key);
	    }
	}
      else if (scm_is_eq (alg, ps384))
	{
	  uint8_t padding[SHA384_DIGEST_SIZE];
	  webid_oidc_random (SHA384_DIGEST_SIZE, padding);
	  if (rsa_pss_sha384_sign_digest_tr
	      (&c_pub, &c_key, NULL, &generate_random, SHA384_DIGEST_SIZE,
	       padding, c_digest, c_sig))
	    {
	      ret = wrap_mpz_t (c_sig);
	    }
	  else
	    {
	      scm_wrong_type_arg ("sign", SCM_ARG2, key);
	    }
	}
      else if (scm_is_eq (alg, ps512))
	{
	  uint8_t padding[SHA512_DIGEST_SIZE];
	  webid_oidc_random (SHA512_DIGEST_SIZE, padding);
	  if (rsa_pss_sha512_sign_digest_tr
	      (&c_pub, &c_key, NULL, &generate_random, SHA512_DIGEST_SIZE,
	       padding, c_digest, c_sig))
	    {
	      ret = wrap_mpz_t (c_sig);
	    }
	  else
	    {
	      scm_wrong_type_arg ("sign", SCM_ARG2, key);
	    }
	}
      else
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      scm_dynwind_end ();
      return ret;
    }
  return SCM_UNDEFINED;
}

SCM_DEFINE (webidoidc_jws_verify_g, "verify", 4, 0, 0,
	    (SCM alg, SCM key, SCM payload, SCM signature),
	    "Verify that @var{signature} is the signature of @var{payload} with @var{key}, using the given algorithm @var{alg}. If @var{alg} is not @code{'RS256}, @code{'RS512}, @code{'ES256}, @code{'ES384}, @code{'ES512}, @code{'PS256}, @code{'PS384}, or @code{'PS512}, or @var{alg} cannot be used with @var{key}, throw @code{incompatible-alg}.")
{
  SCM digest;
  size_t c_digest_size;
  uint8_t *c_digest;
  if (scm_is_eq (alg, rs256)
      || scm_is_eq (alg, es256) || scm_is_eq (alg, ps256))
    {
      digest = webidoidc_hash_g (sha256, payload);
    }
  else if (scm_is_eq (alg, es384) || scm_is_eq (alg, ps384))
    {
      digest = webidoidc_hash_g (sha384, payload);
    }
  else if (scm_is_eq (alg, rs512)
	   || scm_is_eq (alg, es512) || scm_is_eq (alg, ps512))
    {
      digest = webidoidc_hash_g (sha512, payload);
    }
  else
    {
      scm_throw (incompatible_alg, scm_list_2 (alg, key));
    }
  c_digest = get_as_bytevector (digest, &c_digest_size, 1);
  if (scm_is_eq (alg, es256) || scm_is_eq (alg, es384)
      || scm_is_eq (alg, es512))
    {
      struct ecc_point c_pub;
      SCM crv = scm_assq_ref (key, kcrv);
      const struct ecc_curve *c_crv = do_ecc_curve_load (crv, 1);
      struct dsa_signature sig;
      size_t c_sig_size;
      uint8_t *c_sig;
      SCM r, s;
      int ok = 0;
      c_sig = get_as_bytevector (signature, &c_sig_size, 1);
      r = wrap_bytevector (c_sig_size / 2, c_sig);
      s =
	wrap_bytevector (c_sig_size - c_sig_size / 2,
			 c_sig + (c_sig_size / 2));
      scm_dynwind_begin (0);
      ecc_point_init (&c_pub, c_crv);
      dynwind_ecc_point_clear (&c_pub);
      dsa_signature_init (&sig);
      dynwind_dsa_signature_clear (&sig);
      if (!do_ecc_point_load (&c_pub, key))
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      do_mpz_t_load (sig.r, r, 1);
      do_mpz_t_load (sig.s, s, 1);
      ok = ecdsa_verify (&c_pub, c_digest_size, c_digest, &sig);
      scm_dynwind_end ();
      return scm_from_bool (ok);
    }
  else
    {
      struct rsa_public_key c_pub;
      mpz_t c_sig;
      int ok = 0;
      scm_dynwind_begin (0);
      rsa_public_key_init (&c_pub);
      dynwind_rsa_public_key_clear (&c_pub);
      if (!do_rsa_public_key_load (&c_pub, key))
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      mpz_init (c_sig);
      dynwind_mpz_t_clear (&c_sig);
      do_mpz_t_load (c_sig, signature, 1);
      if (scm_is_eq (alg, rs256))
	{
	  ok = rsa_sha256_verify_digest (&c_pub, c_digest, c_sig);
	}
      else if (scm_is_eq (alg, rs512))
	{
	  ok = rsa_sha512_verify_digest (&c_pub, c_digest, c_sig);
	}
      else if (scm_is_eq (alg, ps256))
	{
	  ok =
	    rsa_pss_sha256_verify_digest (&c_pub, c_digest_size, c_digest,
					  c_sig);
	}
      else if (scm_is_eq (alg, ps384))
	{
	  ok =
	    rsa_pss_sha384_verify_digest (&c_pub, c_digest_size, c_digest,
					  c_sig);
	}
      else if (scm_is_eq (alg, ps512))
	{
	  ok =
	    rsa_pss_sha512_verify_digest (&c_pub, c_digest_size, c_digest,
					  c_sig);
	}
      else
	{
	  scm_throw (incompatible_alg, scm_list_2 (alg, key));
	}
      scm_dynwind_end ();
      return scm_from_bool (ok);
    }
}

void
init_webidoidc_jws (void)
{
#ifndef SCM_MAGIC_SNARFER
#include "libwebidoidc-jws.x"
#endif /* not SCM_MAGIC_SNARFER */
}