;; 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 jws)
  #:use-module (webid-oidc jwk)
  #:use-module (webid-oidc errors)
  #:use-module (webid-oidc web-i18n)
  #:use-module ((webid-oidc stubs) #:prefix stubs:)
  #:use-module (rnrs bytevectors)
  #:use-module (ice-9 receive)
  #:use-module (ice-9 exceptions)
  #:use-module (ice-9 match)
  #:use-module (oop goops)
  #:declarative? #t
  #:export
  (

   &invalid-jws
   make-invalid-jws
   invalid-jws?

   the-jws
   jws?

   jws-alg

   &cannot-query-identity-provider
   make-cannot-query-identity-provider
   cannot-query-identity-provider?
   cannot-query-identity-provider-value

   &signed-in-future
   make-signed-in-future
   signed-in-future?
   error-signature-date
   error-current-date

   &expired
   make-expired
   expired?
   error-expiration-date
   ;; error-current-date works for that one too

   jws-decode
   jws-encode

   ))

(define-exception-type
  &invalid-jws
  &external-error
  make-invalid-jws
  invalid-jws?)

(define (the-jws x)
  (with-exception-handler
      (lambda (error)
        (let ((final-message
               (if (exception-with-message? error)
                   (format #f (G_ "the JWS is invalid: ~a")
                           (exception-message error))
                   (format #f (G_ "the JWS is invalid")))))
          (raise-exception
           (make-exception
            (make-invalid-jws)
            (make-exception-with-message final-message)
            error))))
    (lambda ()
      (match x
        ((header . payload)
         (let examine-header ((header header)
                              (alg #f)
                              (other-header-fields '()))
           (match header
             (()
              (let examine-payload ((payload payload)
                                    (other-payload-fields '()))
                (match payload
                  (()
                   (unless alg
                     (fail (format #f (G_ "the JWS header does not have an \"alg\" field"))))
                   `(((alg . ,(symbol->string alg))
                      ,@(reverse other-header-fields))
                     . ,(reverse other-payload-fields)))
                  ((((? symbol? key) . value) payload ...)
                   (examine-payload payload
                                    `((,key . ,value) ,@other-payload-fields)))
                  (else
                   (fail (format #f (G_ "invalid JSON object as payload")))))))
             ((('alg . (? string? given-alg)) header ...)
              (case (string->symbol given-alg)
                ((HS256 HS384 HS512
                        RS256 RS384 RS512
                        ES256 ES384 ES512
                        PS256 PS384 PS512)
                 #t)
                (else
                 (fail (format #f (G_ "invalid signature algorithm: ~s") given-alg))))
              (examine-header header (or alg (string->symbol given-alg))
                              other-header-fields))
             ((('alg . invalid) header ...)
              (fail (format #f (G_ "invalid \"alg\" value: ~s") invalid)))
             ((((? symbol? key) . value) header ...)
              (examine-header header alg
                              `((,key . ,value) ,@other-header-fields)))
             (else
              (fail (format #f (G_ "invalid JSON object as header")))))))
        (else
         (fail (format #f (G_ "this is not a pair"))))))))

(define (jws? x)
  (false-if-exception
   (the-jws x)))

(define (jws-alg jws)
  (match (the-jws jws)
    ((header . _)
     (string->symbol (assq-ref header 'alg)))))

(define (split-in-3-parts string separator)
  (match (string-split string separator)
    ((header payload signature)
     (values header payload signature))
    (else
     (let ((final-message
            (format #f (G_ "the encoded JWS is not in 3 parts"))))
       (raise-exception
        (make-exception
         (make-invalid-jws)
         (make-exception-with-message final-message)))))))

(define (base64-decode-json str)
  (with-exception-handler
      (lambda (error)
        (let ((final-message
               (if (exception-with-message? error)
                   (format #f (G_ "the encoded JWS header or payload is not a JSON object encoded in base64: ~a")
                           (exception-message error))
                   (format #f (G_ "the encoded JWS header or payload is not a JSON object encoded in base64")))))
          (raise-exception
           (make-exception
            (make-invalid-jws)
            (make-exception-with-message final-message)
            error))))
    (lambda ()
      (stubs:json-string->scm
       (utf8->string (stubs:base64-decode str))))))

(define-exception-type
  &cannot-query-identity-provider
  &external-error
  make-cannot-query-identity-provider
  cannot-query-identity-provider?
  (identity-provider cannot-query-identity-provider-value))

(define-exception-type
  &signed-in-future
  &external-error
  make-signed-in-future
  signed-in-future?
  (signature-date error-signature-date)
  (current-date error-current-date*))

(define-exception-type
  &expired
  &external-error
  make-expired
  expired?
  (expiration-date error-expiration-date)
  (current-date error-current-date**))

(define error-current-date
  (match-lambda
    ((or ($ &signed-in-future _ date)
         ($ &expired _ date)
         ($ &compound-exception (($ &signed-in-future _ date) _ ...))
         ($ &compound-exception (($ &expired _ date) _ ...)))
     date)
    (($ &compound-exception (_ sub-exceptions ...))
     (error-current-date (apply make-exception sub-exceptions)))
    (else #f)))

(define (parse str verify)
  (receive (header payload signature)
      (split-in-3-parts str #\.)
    (let ((base (string-append header "." payload))
          (header (base64-decode-json header))
          (payload (base64-decode-json payload)))
      (let ((ret `(,header . ,payload)))
        (verify ret base signature)
        ret))))

(define (verify-any alg keys payload signature)
  (let try-with-key ((keys keys))
    (match keys
      (()
       (let ((final-message
              (format #f (G_ "the JWS is not signed by any of the expected set of public keys"))))
         (raise-exception
          (make-exception
           (make-invalid-jws)
           (make-exception-with-message final-message)))))
      ((next-key keys ...)
       (with-exception-handler
           (lambda (error)
             (unless (stubs:invalid-signature? error)
               (let ((final-message
                      (if (exception-with-message? error)
                          (format #f (G_ "while verifying the JWS signature: ~a")
                                  (exception-message error))
                          (format #f (G_ "an unexpected error happened while verifying a JWS")))))
                 (raise-exception
                  (make-exception
                   (make-invalid-jws)
                   (make-exception-with-message final-message)
                   error))))
             (try-with-key keys))
         (lambda ()
           (stubs:verify alg (key->jwk next-key) payload signature))
         #:unwind? #t
         #:unwind-for-type stubs:&invalid-signature)))))

;; For verification, we can supply a JWKS, or a public key, or a list
;; of public keys. The JWKS case is handled in (webid-oidc jwk).

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

(define-method (keys (key <key-pair>))
  (list (public-key key)))

(define-method (keys (keys <list>))
  (map public-key keys))

(define (jws-decode str lookup-keys)
  (with-exception-handler
      (lambda (error)
        (let ((final-message
               (if (exception-with-message? error)
                   (format #f (G_ "cannot decode a JWS: ~a")
                           (exception-message error))
                   (format #f (G_ "cannot decode a JWS")))))
          (raise-exception
           (make-exception
            (make-invalid-jws)
            (make-exception-with-message final-message)
            error))))
    (lambda ()
      (parse str
             (lambda (jws payload signature)
               (let ((k (keys (lookup-keys jws))))
                 (verify-any (jws-alg jws) k payload signature)))))))

(define (jws-encode jws key)
  (with-exception-handler
      (lambda (error)
        (let ((final-message
               (if (exception-with-message? error)
                   (format #f (G_ "cannot encode a JWS: ~a")
                           (exception-message error))
                   (format #f (G_ "cannot encode a JWS")))))
          (raise-exception
           (make-exception
            (make-invalid-jws)
            (make-exception-with-message final-message)
            error))))
    (lambda ()
      (match jws
        ((header . payload)
         (let ((header (stubs:scm->json-string header))
               (payload (stubs:scm->json-string payload)))
           (let ((header (stubs:base64-encode header))
                 (payload (stubs:base64-encode payload)))
             (let ((payload (string-append header "." payload)))
               (let ((signature (stubs:sign (jws-alg jws) (key->jwk key) payload)))
                 (string-append payload "." signature))))))))))