;; 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 jti)
  #:use-module (webid-oidc oidc-configuration)
  #:use-module (webid-oidc serializable)
  #:use-module ((webid-oidc stubs) #:prefix stubs:)
  #:use-module ((webid-oidc parameters) #:prefix p:)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-19)
  #:use-module (ice-9 receive)
  #:use-module (ice-9 exceptions)
  #:use-module (ice-9 match)
  #:use-module (ice-9 optargs)
  #:use-module (web uri)
  #:use-module (sxml match)
  #:use-module (oop goops)
  #:declarative? #t
  #:re-export
  (
   (&jti-found . &nonce-found)
   (make-jti-found . make-nonce-found)
   (jti-found? . nonce-found?)
   (jti-found-jti . nonce-found-nonce)
   )
  #:replace
  (
   exp ;; This is a function in guile
   )
  #:export
  (

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

   <token>

   <time-bound-token> iat default-validity has-explicit-exp?
   nonce-field-name ;; DPoP proofs use 'jti instead of 'nonce

   <oidc-token> iss

   <single-use-token> nonce
   
   token->jwt

   &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

   lookup-keys
   verify
   decode
   encode
   issue

   ))

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

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

(define (key-alg key)
  (alg key))

(define-method (initialize (token <token>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((alg #f)
    (signing-key #f)
    (jwt-header #f)
    (jwt-payload #f))
   (let do-initialize ((alg alg)
                       (signing-key signing-key)
                       (jwt-header jwt-header)
                       (jwt-payload jwt-payload))
     (cond
      ((string? alg)
       (do-initialize (string->symbol alg) signing-key jwt-header jwt-payload))
      (alg
       (case alg
         ((HS256 HS384 HS512
                 RS256 RS384 RS512
                 ES256 ES384 ES512
                 PS256 PS384 PS512)
          (slot-set! token 'alg alg))
         (else
          (raise-exception
           (make-exception
            (make-invalid-jws)
            (make-exception-with-message
             (format #f (G_ "unsupported JWS algorithm: ~s") alg)))))))
      (signing-key
       (do-initialize (key-alg signing-key) #f jwt-payload jwt-header))
      ((and jwt-header jwt-payload)
       (do-initialize (assq-ref jwt-header 'alg) #f #f #f))
      (else
       (raise-exception
        (make-exception
         (make-invalid-jws)
         (make-exception-with-message
          (G_ "when making a token either #:alg or (#:jwt-header and #:jwt-payload) should be passed")))))))))

(define-class <generic-with-default> (<generic>)
  ;; neutral is the list of values that are returned when there are no
  ;; next methods.
  (neutral #:init-keyword #:neutral))

(define-method (no-next-method (generic <generic-with-default>) args)
  (apply values (slot-ref generic 'neutral)))

(define-method (no-applicable-method (generic <generic-with-default>) args)
  (apply values (slot-ref generic 'neutral)))

(define (date->sxml date)
  (number->string (time-second (date->time-utc date))))

(define-class <time-bound-token> (<token>)
  (iat #:init-keyword #:iat #:accessor iat #:->sxml date->sxml)
  (exp #:init-keyword #:exp #:accessor exp #:->sxml date->sxml)
  #:module-name '(webid-oidc jws))

(define default-validity
  (make <generic-with-default>
    #:name 'default-validity
    #:neutral (list #f)))

(define-method (has-explicit-exp? (token <time-bound-token>))
  ;; Change it to #f when the token should not have an explicit
  ;; expiration date, such as DPoP proofs
  #t)

(define-method (initialize (token <time-bound-token>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((iat ((p:current-date)))
    (exp #f)
    (validity (default-validity token))
    (jwt-header #f)
    (jwt-payload #f))
   (let do-initialize ((iat iat)
                       (exp exp)
                       (validity validity)
                       (jwt-header jwt-header)
                       (jwt-payload jwt-payload))
     (cond
      ((string? iat)
       (do-initialize (string->number iat) exp validity jwt-header jwt-payload))
      ((integer? iat)
       (do-initialize (make-time time-utc 0 iat) exp validity jwt-header jwt-payload))
      ((time? iat)
       (do-initialize (time-utc->date iat) exp validity jwt-header jwt-payload))
      ((and (not exp) (date? iat) (integer? validity))
       (do-initialize iat
                      (+ (time-second (date->time-utc iat))
                         validity)
                      validity
                      jwt-header
                      jwt-payload))
      ((string? exp)
       (do-initialize iat (string->number exp) validity jwt-header jwt-payload))
      ((integer? exp)
       (do-initialize iat (make-time time-utc 0 exp) validity jwt-header jwt-payload))
      ((time? exp)
       (do-initialize iat (time-utc->date exp) validity jwt-header jwt-payload))
      ((and jwt-header jwt-payload)
       (do-initialize (assq-ref jwt-payload 'iat)
                      (and (has-explicit-exp? token)
                           (assq-ref jwt-payload 'exp))
                      validity #f #f))
      ((and iat exp)
       (unless (date? iat)
         (scm-error 'wrong-type-arg "make"
                    (G_ "#:iat should be a date")
                    '()
                    (list iat)))
       (unless (date? exp)
         (scm-error 'wrong-type-arg "make"
                    (G_ "#:exp should be a date")
                    '()
                    (list exp)))
       (slot-set! token 'iat iat)
       (slot-set! token 'exp exp))
      (else
       (raise-exception
        (make-exception
         (make-invalid-jws)
         (make-exception-with-message
          (G_ "when making a time-bound token, either its required fields (#:iat, and either #:exp or #:validity) or (#:jwt-header and #:jwt-payload) should be passed")))))))))

(define-class <oidc-token> (<token>)
  (iss #:init-keyword #:iss #:accessor iss #:->sxml uri->string)
  #:module-name '(webid-oidc jws))

(define-method (default-validity (token <oidc-token>))
  (let ((next (next-method))
        (mine (p:oidc-token-default-validity)))
    (if (and next (< next mine))
        next
        mine)))

(define-method (initialize (token <oidc-token>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((iss #f)
    (jwt-header #f)
    (jwt-payload #f))
   (let do-initialize ((iss iss)
                       (jwt-header jwt-header)
                       (jwt-payload jwt-payload))
     (cond
      ((string? iss)
       (do-initialize (string->uri iss) jwt-header jwt-payload))
      (iss
       (unless (uri? iss)
         (scm-error 'wrong-type-arg "make"
                    (G_ "#:iss should be an URI")
                    '()
                    (list iss)))
       (slot-set! token 'iss iss))
      ((and jwt-header jwt-payload)
       (do-initialize (assq-ref jwt-payload 'iss) #f #f))
      (else
       (raise-exception
        (make-exception
         (make-invalid-jws)
         (make-exception-with-message
          (G_ "when making an OIDC token, either its required #:iss field or (#:jwt-header and #:jwt-payload) should be passed")))))))))

(define-class <single-use-token> (<time-bound-token>)
  (nonce #:init-keyword #:nonce #:accessor nonce)
  #:module-name '(webid-oidc jws))

(define-method (default-validity (token <single-use-token>))
  (let ((next (next-method))
        (mine (p:authorization-code-default-validity)))
    (if (and next (< next mine))
        next
        mine)))

(define nonce-field-name
  (make <generic-with-default>
    #:name 'nonce-field-name
    #:neutral (list 'nonce)))

(define-method (nonce-field-name (token <top>))
  ;; Without this method, this is an infinite loop.
  (next-method))

(define-method (initialize (token <single-use-token>) initargs)
  (next-method)
  (let-keywords
   initargs #t
   ((nonce (stubs:random 12))
    (jwt-header #f)
    (jwt-payload #f))
   ;; The maximum validity is 2 minutes
   (let ((iat (time-second (date->time-utc (iat token))))
         (exp (time-second (date->time-utc (exp token)))))
     (let ((validity (- exp iat)))
       (when (> validity 120)
         (let ((true-exp (+ iat 120)))
           (slot-set! token 'exp (time-utc->date (make-time time-utc 0 true-exp)))))))
   (let do-initialize ((nonce nonce)
                       (jwt-header jwt-header)
                       (jwt-payload jwt-payload))
     (cond
      ((and jwt-header jwt-payload)
       (do-initialize (assq-ref jwt-payload (nonce-field-name token)) #f #f))
      (nonce
       (unless (string? nonce)
         (scm-error 'wrong-type-arg "make"
                    (G_ "#:nonce should be a string")
                    '()
                    (list nonce)))
       (slot-set! token 'nonce nonce))
      (else
       (raise-exception
        (make-exception
         (make-invalid-jws)
         (make-exception-with-message
          (G_ "when making a single-use token, either its required #:nonce field or (#:jwt-header and #:jwt-payload) should be passed")))))))))

(define token->jwt
  (make <generic-with-default>
    #:name 'token->jwt
    #:neutral (list '() '())))

(define-method (token->jwt (token <token>))
  (receive (base-header base-payload)
      (next-method)
    (values
     `((alg . ,(symbol->string (alg token)))
       ,@base-header)
     base-payload)))

(define-method (token->jwt (token <time-bound-token>))
  (receive (base-header base-payload)
      (next-method)
    (values base-header
            `((iat . ,(time-second (date->time-utc (iat token))))
              ,@(if (has-explicit-exp? token)
                    `((exp . ,(time-second (date->time-utc (exp token)))))
                    '())
              ,@base-payload))))

(define-method (token->jwt (token <single-use-token>))
  (receive (base-header base-payload)
      (next-method)
    (values base-header
            `((,(nonce-field-name token) . ,(nonce token))
              ,@base-payload))))

(define-method (token->jwt (token <oidc-token>))
  (receive (base-header base-payload)
      (next-method)
    (values base-header
            `((iss . ,(uri->string (iss token)))
              ,@base-payload))))

(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 token-class str verify-signature)
  (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 (make token-class #:jwt-header header #:jwt-payload payload)))
        (verify-signature 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 lookup-keys
  (make <generic-with-default>
    #:name 'lookup-keys
    #:neutral (list '())))

(define-method (lookup-keys (token <oidc-token>) args)
  (let ((iss (iss token)))
    (let ((cfg
           (with-exception-handler
               (lambda (error)
                 (let ((final-message
                        (if (exception-with-message? error)
                            (format #f (G_ "I cannot query the identity provider configuration: ~a")
                                    (exception-message error))
                            (format #f (G_ "I cannot query the identity provider configuration")))))
                   (raise-exception
                    (make-exception
                     (make-cannot-query-identity-provider iss)
                     (make-exception-with-message final-message)
                     error))))
             (lambda ()
               (make <oidc-configuration>
                 #:server iss)))))
      (with-exception-handler
          (lambda (error)
            (raise-exception
             (make-exception
              (make-cannot-query-identity-provider iss)
              (make-exception-with-message
               (if (exception-with-message? error)
                   (format #f (G_ "I cannot query the JWKS URI of the identity provider: ~a")
                           (exception-message error))
                   (format #f (G_ "I cannot query the JWKS URI of the identity provider")))))))
        (lambda ()
          (append
           (keys (next-method))
           (keys (jwks cfg))))))))

(define verify
  (make <generic-with-default>
    #:name 'verify
    #:neutral (list #t)))

(define-method (verify (token <time-bound-token>) args)
  (next-method)
  (let-keywords
   args #t
   ((current-date ((p:current-date))))
   (let ((iat (iat token))
         (exp (exp token)))
     (let ((iat-s (time-second (date->time-utc iat)))
           (exp-s (time-second (date->time-utc exp)))
           (current-s (time-second (date->time-utc current-date))))
       (when (>= iat-s (+ current-s 5))
         (let ((final-message
                (format #f (G_ "the token is signed in the future, ~a, relative to current ~a")
                        (date->string iat)
                        (date->string current-date))))
           (raise-exception
            (make-exception
             (make-signed-in-future iat current-date)
             (make-exception-with-message final-message)))))
       (when (>= current-s exp-s)
         (let ((final-message
                (format #f (G_ "the token expired ~a, which is in the past (from ~a)")
                        (date->string exp)
                        (date->string current-date))))
           (raise-exception
            (make-exception
             (make-expired exp current-date)
             (make-exception-with-message final-message)))))))))

(define-method (verify (token <single-use-token>) args)
  (next-method)
  (let-keywords
   args #t
   ((current-date ((p:current-date))))
   (let ((exp (exp token)))
     (let ((exp-s (time-second (date->time-utc exp)))
           (current-s (time-second (date->time-utc current-date))))
       (jti-check (nonce token) (- exp-s current-s))))))

(define* (decode token-class str . args)
  (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 token-class str
             (lambda (token payload signature)
               (let ((k (keys (lookup-keys token args))))
                 (verify-any (alg token) k payload signature))
               (verify token args))))))

(define (encode token 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 ()
      (receive (header payload) (token->jwt token)
        (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 (alg token) (key->jwk key) payload)))
                (string-append payload "." signature)))))))))

(define* (issue token-class issuer-key . args)
  (encode (apply make token-class #:signing-key issuer-key args) issuer-key))