;; disfluid, 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/>.

(define-module (webid-oidc jwk)
  #:use-module ((webid-oidc stubs) #:prefix stubs:)
  #:use-module ((webid-oidc parameters) #:prefix p:)
  #:use-module (webid-oidc errors)
  #:use-module (webid-oidc web-i18n)
  #:use-module (webid-oidc serializable)
  #:use-module (ice-9 receive)
  #:use-module (ice-9 optargs)
  #:use-module (ice-9 exceptions)
  #:use-module (ice-9 hash-table)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-19)
  #:use-module (web response)
  #:use-module (rnrs bytevectors)
  #:use-module (oop goops)
  #:use-module (sxml match)
  #:declarative? #t
  #:export
  (
   <private-key> alg
   <public-key>
   <key-pair> public-key private-key
   <rsa-key-pair>
   <ec-key-pair>
   <rsa-private-key> rsa-d rsa-p rsa-q rsa-dp rsa-dq rsa-qi
   <rsa-public-key> rsa-n rsa-e
   <ec-scalar> ec-crv ec-x ec-y
   <ec-point> ec-z
   <jwks> keys

   check-key
   key->jwk
   jwk->key
   kty
   jkt
   generate-key
   serve
   get-jwks

   &not-a-jwk
   make-not-a-jwk
   not-a-jwk?

   &not-a-jwks
   make-not-a-jwks
   not-a-jwks?
   ))

(define-exception-type
  &not-a-jwk
  &external-error
  make-not-a-jwk
  not-a-jwk?)

(define-exception-type
  &not-a-jwks
  &external-error
  make-not-a-jwks
  not-a-jwks?)

(define-class <private-key> ()
  (alg #:init-keyword #:alg #:accessor alg)
  #:metaclass <plugin-class>
  #:module-name '(webid-oidc jwk))

(define-class <public-key> ()
  #:metaclass <plugin-class>
  #:module-name '(webid-oidc jwk))

(define-class <key-pair> ()
  (public-key #:init-keyword #:public-key #:accessor public-key)
  (private-key #:init-keyword #:private-key #:accessor private-key)
  #:metaclass <plugin-class>
  #:module-name '(webid-oidc jwk))

(define-class <rsa-key-pair> (<key-pair>)
  #:module-name '(webid-oidc jwk))

(define-class <ec-key-pair> (<key-pair>)
  (crv #:init-keyword #:crv #:accessor ec-crv)
  #:module-name '(webid-oidc jwk))

(define-class <rsa-private-key> (<private-key>)
  (d #:init-keyword #:d #:accessor rsa-d)
  (p #:init-keyword #:p #:accessor rsa-p)
  (q #:init-keyword #:q #:accessor rsa-q)
  (dp #:init-keyword #:dp #:accessor rsa-dp)
  (dq #:init-keyword #:dq #:accessor rsa-dq)
  (qi #:init-keyword #:qi #:accessor rsa-qi)
  #:module-name '(webid-oidc jwk))

(define-class <rsa-public-key> (<public-key>)
  (n #:init-keyword #:n #:accessor rsa-n)
  (e #:init-keyword #:e #:accessor rsa-e)
  #:module-name '(webid-oidc jwk))

(define-class <ec-scalar> (<private-key>)
  (crv #:init-keyword #:crv #:accessor ec-crv)
  (z #:init-keyword #:z #:accessor ec-z)
  #:module-name '(webid-oidc jwk))

(define-class <ec-point> (<public-key>)
  (crv #:init-keyword #:crv #:accessor ec-crv)
  (x #:init-keyword #:x #:accessor ec-x)
  (y #:init-keyword #:y #:accessor ec-y)
  #:module-name '(webid-oidc jwk))

(define-method (initialize-key-pair (key <key-pair>) (public <rsa-public-key>) (private <rsa-private-key>))
  (set! (public-key key) public)
  (set! (private-key key) private))

(define-method (initialize-key-pair (key <key-pair>) (public <ec-point>) (private <ec-scalar>))
  (set! (public-key key) public)
  (set! (private-key key) private))

(define-method (initialize (key <key-pair>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((public-key #f)
    (private-key #f))
   (initialize-key-pair key public-key private-key))
  (check-key key))

(define-method (initialize-rsa-key-pair (key <rsa-key-pair>) (public <rsa-public-key>) (private <rsa-private-key>))
  #t)

(define-method (initialize (key <rsa-key-pair>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((public-key #f)
    (private-key #f))
   (initialize-rsa-key-pair key public-key private-key))
  (check-key key))

(define-method (initialize-ec-key-pair (key <ec-key-pair>) (public <ec-point>) (private <ec-scalar>))
  (unless (eq? (ec-crv public) (ec-crv private))
    (raise-exception
     (make-exception
      (make-not-a-jwk)
      (make-exception-with-message (G_ "the point and scalar are not on the same curve")))))
  (set! (ec-crv key) (ec-crv public)))

(define-method (initialize (key <ec-key-pair>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((public-key #f)
    (private-key #f))
   (initialize-ec-key-pair key public-key private-key)
   (check-key key)))

(define-method (initialize (key <rsa-private-key>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((alg #f))
   (when (string? alg)
     (set! alg (string->symbol alg)))
   (slot-set! key 'alg (or alg 'RS256)))
  (check-key key))

(define-method (initialize (key <rsa-public-key>) initargs)
  (next-method)
  (check-key key))

(define-method (initialize (key <ec-point>) initargs)
  (next-method)
  (when (string? (ec-crv key))
    (set! (ec-crv key) (string->symbol (ec-crv key))))
  (check-key key))

(define-method (initialize (key <ec-scalar>) initargs)
  (next-method)
  (when (string? (ec-crv key))
    (set! (ec-crv key) (string->symbol (ec-crv key))))
  (let-keywords
   initargs #t
   ((alg #f))
   (when (string? alg)
     (set! alg (string->symbol alg)))
   (slot-set! key 'alg (or alg 'ES256)))
  (check-key key))

(define-method (alg (key <key-pair>))
  (alg (private-key key)))

(define-method (rsa-d (key <rsa-key-pair>))
  (rsa-d (private-key key)))

(define-method (rsa-p (key <rsa-key-pair>))
  (rsa-p (private-key key)))

(define-method (rsa-q (key <rsa-key-pair>))
  (rsa-q (private-key key)))

(define-method (rsa-dp (key <rsa-key-pair>))
  (rsa-dp (private-key key)))

(define-method (rsa-dq (key <rsa-key-pair>))
  (rsa-dq (private-key key)))

(define-method (rsa-qi (key <rsa-key-pair>))
  (rsa-qi (private-key key)))

(define-method (rsa-n (key <rsa-key-pair>))
  (rsa-n (public-key key)))

(define-method (rsa-e (key <rsa-key-pair>))
  (rsa-e (public-key key)))

(define-method (ec-x (key <ec-key-pair>))
  (ec-x (public-key key)))

(define-method (ec-y (key <ec-key-pair>))
  (ec-y (public-key key)))

(define-method (ec-z (key <ec-key-pair>))
  (ec-z (private-key key)))

(define-method (equal? (x <key-pair>) (y <key-pair>))
  (and (equal? (public-key x) (public-key y))
       (equal? (private-key x) (private-key y))))

(define-method (equal? (x <public-key>) (y <public-key>))
  #f)

(define-method (equal? (x <private-key>) (y <private-key>))
  #f)

(define-method (equal? (x <rsa-public-key>) (y <rsa-public-key>))
  (and (equal? (rsa-n x) (rsa-n y))
       (equal? (rsa-e x) (rsa-e y))))

(define-method (equal? (x <rsa-private-key>) (y <rsa-private-key>))
  (and (equal? (alg x) (alg y))
       (equal? (rsa-d x) (rsa-d y))
       (equal? (rsa-p x) (rsa-p y))
       (equal? (rsa-q x) (rsa-q y))
       (equal? (rsa-dp x) (rsa-dp y))
       (equal? (rsa-dq x) (rsa-dq y))
       (equal? (rsa-qi x) (rsa-qi y))))

(define-method (equal? (x <ec-point>) (y <ec-point>))
  (and (equal? (ec-x x) (ec-x y))
       (equal? (ec-y x) (ec-y y))))

(define-method (equal? (x <ec-scalar>) (y <ec-scalar>))
  (and (equal? (alg x) (alg y))
       (equal? (ec-z x) (ec-z y))))

(define (check-and-kty key)
  (with-exception-handler
      (lambda (error)
        (let ((final-message
               (if (exception-with-message? error)
                   (format #f (G_ "the JWK is invalid: ~a")
                           (exception-message error))
                   (format #f (G_ "the JWK is invalid")))))
          (raise-exception
           (make-exception
            (make-not-a-jwk)
            (make-exception-with-message final-message)
            error))))
    (lambda ()
      (let ((kty (stubs:kty (key->jwk key))))
        (unless kty
          (fail (G_ "cannot compute the key type")))
        kty))))

(define-method (key->jwk (key <key-pair>))
  ;; kty and crv fields are present in both the public and private
  ;; key, but they must not be duplicated, and we want to keep all
  ;; fields in order.
  (let ((with-duplicates
         (append (key->jwk (public-key key))
                 (key->jwk (private-key key)))))
  (let ((lookup
         (alist->hash-table with-duplicates)))
    (let keep-unique-in-order ((order (map car with-duplicates))
                               (constructed '()))
      (match order
        (()
         (reverse constructed))
        ((hd tl ...)
         (let ((found (hash-ref lookup hd)))
           (when found
             (hash-remove! lookup hd))
           (if found
               (keep-unique-in-order tl `((,hd . ,found) ,@constructed))
               (keep-unique-in-order tl constructed)))))))))

(define-method (key->jwk (key <rsa-private-key>))
  `((kty . ,(symbol->string (kty key)))
    (alg . ,(symbol->string (alg key)))
    (d . ,(rsa-d key))
    (p  . ,(rsa-p key))
    (q . ,(rsa-q key))
    (dp . ,(rsa-dp key))
    (dq . ,(rsa-dq key))
    (qi . ,(rsa-qi key))))

(define-method (key->jwk (key <rsa-public-key>))
  `((kty . ,(symbol->string (kty key)))
    (n . ,(rsa-n key))
    (e . ,(rsa-e key))))

(define-method (key->jwk (key <ec-point>))
  `((crv . ,(symbol->string (ec-crv key)))
    (kty . ,(symbol->string (kty key)))
    (x . ,(ec-x key))
    (y . ,(ec-y key))))

(define-method (key->jwk (key <ec-scalar>))
  `((crv . ,(symbol->string (ec-crv key)))
    (kty . ,(symbol->string (kty key)))
    (alg . ,(symbol->string (alg key)))
    (z . ,(ec-z key))))

(define-method (check-key key)
  (check-and-kty (key->jwk key)))

(define (check-rsa-key key)
  (unless (eq? (check-and-kty key) 'RSA)
    (raise-exception
     (make-exception
      (make-not-a-jwk)
      (make-exception-with-message
       (format #f (G_ "it is built as an RSA key or key pair, but it is not")))))))

(define (check-ec-key key)
  (unless (eq? (check-and-kty key) 'EC)
    (raise-exception
     (make-exception
      (make-not-a-jwk)
      (make-exception-with-message
       (format #f (G_ "it is built as an elliptic curve key or key pair, but it is not")))))))

(define-method (check-key (key <rsa-key-pair>))
  (check-rsa-key key))

(define-method (check-key (key <rsa-public-key>))
  (check-rsa-key key))

(define-method (check-key (key <rsa-private-key>))
  (check-rsa-key key))

(define-method (check-key (key <ec-key-pair>))
  (check-ec-key key))

(define-method (check-key (key <ec-point>))
  (check-ec-key key))

(define-method (check-key (key <ec-scalar>))
  (check-ec-key key))

(define-method (kty (key <rsa-key-pair>)) 'RSA)
(define-method (kty (key <rsa-public-key>)) 'RSA)
(define-method (kty (key <rsa-private-key>)) 'RSA)

(define-method (kty (key <ec-key-pair>)) 'EC)
(define-method (kty (key <ec-point>)) 'EC)
(define-method (kty (key <ec-scalar>)) 'EC)

(define-method (public-key (key <public-key>))
  key)

(define-method (private-key (key <private-key>))
  key)

(define (jwk->key fields)
  (let ((kty (stubs:kty fields))
        (alg (assq-ref fields 'alg)))
    (let ((explicit-kty (assq-ref fields 'kty)))
      (when (and kty explicit-kty (not (eq? kty (string->symbol explicit-kty))))
        (raise-exception
         (make-exception
          (make-not-a-jwk)
          (make-exception-with-message (format #f (G_ "the key advertises a key type of ~s, but actually it is ~s")
                                               explicit-kty kty))))))
    (case kty
      ((RSA)
       (let ((d (assq-ref fields 'd))
             (p (assq-ref fields 'p))
             (q (assq-ref fields 'q))
             (dp (assq-ref fields 'dp))
             (dq (assq-ref fields 'dq))
             (qi (assq-ref fields 'qi))
             (n (assq-ref fields 'n))
             (e (assq-ref fields 'e)))
         (let ((public
                (and n e
                     (make <rsa-public-key> #:n n #:e e)))
               (private
                (and d p q dp dq qi
                     (make <rsa-private-key>
                       #:alg (and alg (string->symbol alg))
                       #:d d
                       #:p p
                       #:q q
                       #:dp dp
                       #:dq dq
                       #:qi qi))))
           (if (and public private)
               (make <rsa-key-pair> #:public-key public #:private-key private)
               (or public private)))))
      ((EC)
       (let ((crv (string->symbol (assq-ref fields 'crv)))
             (x (assq-ref fields 'x))
             (y (assq-ref fields 'y))
             (z (assq-ref fields 'z)))
         (let ((public
                (and x y
                     (make <ec-point> #:crv crv #:x x #:y y)))
               (private
                (and z
                     (make <ec-scalar>
                       #:alg (and alg (string->symbol alg))
                       #:crv crv
                       #:z z))))
           (if (and public private)
               (make <ec-key-pair> #:public-key public #:private-key private)
               (or public private)))))
      (else
       (raise-exception
        (make-exception
         (make-not-a-jwk)
         (make-exception-with-message (G_ "this is neither a RSA key nor an elliptic curve key"))))))))

(define (jkt x)
  (stubs:jkt (key->jwk x)))

(define (generate-key . args)
  (jwk->key (apply stubs:generate-key args)))

(define-class <jwks> ()
  (keys #:init-keyword #:keys #:accessor keys))

(define-method (initialize (jwks <jwks>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((keys '()))
   (slot-set! jwks 'keys (map public-key keys))))

(define-method (serve (jwks <jwks>) expiration-date)
  (values
   (build-response
    #:headers `((content-type . (application/json))
                (expires . ,expiration-date)))
   (stubs:scm->json-string
    `((keys
       . ,(list->vector
           (map key->jwk (keys jwks))))))))

(define (get-jwks uri)
  (receive (response response-body) ((p:anonymous-http-request) uri)
    (with-exception-handler
        (lambda (error)
          (raise-exception
           (make-exception
            (make-not-a-jwks)
            (make-exception-with-message
             (if (exception-with-message? error)
                 (format #f (G_ "cannot fetch a JWKS: ~a")
                         (exception-message error))
                 (format #f (G_ "cannot fetch a JWKS"))))
            error)))
      (lambda ()
        (unless (eqv? (response-code response) 200)
          (fail (format #f (G_ "the request failed with ~s ~s")
                        (response-code response)
                        (response-reason-phrase response))))
        (let ((content-type (response-content-type response)))
          (unless content-type
            (fail (format #f (G_ "missing content-type"))))
          (unless (and (eq? (car content-type) 'application/json)
                       (or (equal? (assoc-ref (cdr content-type) 'charset)
                                   "utf-8")
                           (not (assoc-ref (cdr content-type) 'charset))))
            (fail (format #f (G_ "invalid content-type: ~s") content-type)))
          (unless (string? response-body)
            (set! response-body (utf8->string response-body)))
          (let ((data (stubs:json-string->scm response-body)))
            (let ((keys (vector->list (assq-ref data 'keys))))
              (make <jwks> #:keys (map jwk->key keys)))))))))