;;;; Copyright (C) 2020 Julien Lepiller <julien@lepiller.eu>
;;;; 
;;;; This library is free software; you can redistribute it and/or
;;;; modify it under the terms of the GNU Lesser General Public
;;;; License as published by the Free Software Foundation; either
;;;; version 3 of the License, or (at your option) any later version.
;;;; 
;;;; This library 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
;;;; Lesser General Public License for more details.
;;;; 
;;;; You should have received a copy of the GNU Lesser General Public
;;;; License along with this library; if not, write to the Free Software
;;;; Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
;;;; 

(define-module (http-signature asn1)
  #:use-module (gcrypt base64)
  #:use-module (gcrypt pk-crypto)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (rnrs bytevectors)
  #:export (asn1-type?
            make-asn1-type
            asn1-type-name
            asn1-type-type
            asn1-type-primitive?
            asn1-type-constructed?
            asn1-type-parser
            asn1-type-producer
            asn1-type-aggregator

            asn1:bool
            asn1:null
            asn1:int-bv
            asn1:bitstring
            asn1:octetstring
            asn1:object-identifier
            asn1:sequence

            encode-asn1
            decode-asn1))

;; We declare the ASN.1 types here. Each type is associated with its name,
;; its type byte, a parser and a producer.  We will use them to either parse
;; a DER encoded ASN.1 record, or to produce a DER-encoded ASN.1 record.

(define-record-type asn1-type
  (make-asn1-type name type primitive? constructed? parser producer aggregator)
  asn1-type?
  (name asn1-type-name)
  (type asn1-type-type)
  (primitive? asn1-type-primitive?)
  (constructed? asn1-type-constructed?)
  (parser asn1-type-parser)
  (producer asn1-type-producer)
  (aggregator asn1-type-aggregator))

(define asn1:bool
  (make-asn1-type
    "bool" 1 #t #f
    (lambda (bv pos len)
      (unless (equal? len 1)
        (throw 'invalid-bool-value))
      (let ((val (bytevector-u8-ref bv pos)))
        (not (equal? val 0))))
    (lambda (value)
      (let ((bv (make-bytevector 1)))
        (bytevector-u8-set! bv 0 (if value #xff #x00))
        bv))
    (lambda _ (throw 'primitive-type))))

(define asn1:null
  (make-asn1-type
    "null" 5 #t #f
    (lambda (bv pos len)
      (unless (equal? len 0)
        (throw 'invalid-null-value))
      #nil)
    (lambda (value)
      (make-bytevector 0))
    (lambda _ (throw 'primitive-type))))

;(define asn1:int-int
;  (make-asn1-type
;    "int" 2
;    (lambda (bv pos len)
;      (let ((ans (make-bytevector len)))
;        (bytevector-copy! bv pos ans 0 len)
;        (let loop ((num 0) (off 0))
;          (if (> off len)
;            num
;            (loop (+ (* num 256) (bytevector-u8-ref bv off))
;                  (+ off 1))))))
;    (lambda (value)
;      ())))

(define asn1:int-bv
  (make-asn1-type
    "int" 2 #t #f
    (lambda (bv pos len)
      (let ((ans (make-bytevector len)))
        (bytevector-copy! bv pos ans 0 len)
        ans))
    (lambda (value)
      value)
    (lambda _ (throw 'primitive-type))))

;(define asn1:enum asn1:int-int)

(define asn1:bitstring
  (make-asn1-type
    "bitstring" 3 #t #t
    (lambda (bv pos len)
      (let ((ans (make-bytevector (- len 1))))
        (bytevector-copy! bv (+ pos 1) ans 0 (- len 1))
        ans))
    (lambda (value)
      (let ((ans (make-bytevector (+ (bytevector-length value) 1))))
        (bytevector-u8-set! ans 0 0)
        (bytevector-copy! value 0 ans 1 (bytevector-length value))
        ans))
    (lambda (value)
      (apply bv-append value))))

(define asn1:octetstring
  (make-asn1-type
    "octetstring" 4 #t #t
    (lambda (bv pos len)
      (let ((ans (make-bytevector len)))
        (bytevector-copy! bv pos ans 0 len)
        ans))
    (lambda (value)
      value)
    (lambda (value)
      (apply bv-append value))))

(define asn1:object-identifier
  (make-asn1-type
    "object-identifier" 6 #t #f
    (lambda (bv pos len)
      (when (= len 0)
        (throw 'invalid-object-identifier))
      (let loop ((len len) (pos pos) (value '()) (new? #t))
        (if (<= len 0)
            (if (= len 0)
                (let* ((vals (reverse value))
                       (first-identifier (car vals)))
                  (append
                    (cond
                      ((< first-identifier 40) (list 0 (modulo first-identifier 40)))
                      ((< first-identifier 80) (list 1 (modulo first-identifier 40)))
                      (else (list 2 (- first-identifier 80))))
                    (cdr vals)))
                (throw 'value-too-long))
            (let* ((val (bytevector-u8-ref bv pos))
                   (cont? (> val #x80))
                   (val (logand #x7f val)))
              (if new?
                (loop (- len 1) (+ pos 1)
                      (cons val value)
                      (not cont?))
                (loop (- len 1) (+ pos 1)
                      (cons (+ val (* (car value) 128)) (cdr value))
                      (not cont?)))))))
    (lambda (value)
      (match value
        ((x y rest ...)
         (let ((first-byte (+ (* x 40) y)))
           (apply bv-append
             (map
               (lambda (identifier)
                 (define (id-len identifier)
                   (if (< identifier 128)
                       1
                       (+ 1 (id-len (quotient identifier 128)))))
                 (let ((bv (make-bytevector (id-len identifier))))
                   (let loop ((pos (- (bytevector-length bv) 1)) (id identifier))
                     (if (= pos (- (bytevector-length bv) 1))
                       (begin
                         (bytevector-u8-set! bv pos (modulo id 128))
                         (loop (- pos 1) (quotient id 128)))
                       (if (= pos -1)
                           bv
                           (begin
                             (bytevector-u8-set! bv pos (+ #x80 (modulo id 128)))
                             (loop (- pos 1) (quotient id 128))))))))
               (cons first-byte rest)))))
        (_ (throw 'invalid-identifier))))
    (lambda _ (throw 'primitive-type))))

;; Note that a sequence depends on the types of its contents
(define* (asn1:sequence . types)
  (make-asn1-type
    "sequence" #x10 #f #t
    (lambda (bv pos len)
      (let loop ((result '()) (types types) (rest-len len) (pos pos))
        (match types
          (()
           (unless (equal? rest-len 0)
             (throw 'invalid-sequence-type))
           (reverse result))
          ((type types ...)
           (match (decode-asn1-aux bv pos type)
             ((value . npos)
              (loop (cons value result) types (- rest-len (- npos pos)) npos)))))))
    (lambda (value)
      (unless (equal? (length value) (length types))
        (throw 'incompatible-type-and-value-sizes))
      (apply
        bv-append
        (map encode-asn1 types value)))
    (lambda (value) value)))


;; Now these functions are related to parsing or producing DER-encoded values
;; as bytevectors.

(define (decode-asn1-aux bv pos type)
  "Return a value that corresponds to the DER-encoded value of @var{type}
in @var{bv} at @var{pos}.  The return value is a pair of the decoded
value and the new position after reading the bytevector."
  (define (decode-len bv pos)
    (let ((first-byte (bytevector-u8-ref bv pos)))
      (if (> first-byte 127)
        (let ((num-bytes (- first-byte 128)))
          (let loop ((pos (+ pos 1)) (num num-bytes) (result 0))
            (if (eq? num 0)
                result
                (loop (+ pos 1) (- num 1)
                      (+ (bytevector-u8-ref bv pos) (* 256 result))))))
        first-byte)))

  (define (lenlen bv pos)
    (let ((first-byte (bytevector-u8-ref bv pos)))
      (if (> first-byte 127)
          (- first-byte 127); one for the size byte too, so 127 not 128
          1)))

  (define (constructed? type-byte)
    (not (eq? (logand type-byte #x20) 0)))

  (define (decode-asn1-array bv pos type)
    (let loop ((value '()) (pos pos))
      (if (= (bytevector-u8-ref bv pos) 0)
          (begin
            (unless (= (bytevector-u8-ref bv (+ pos 1)) 0)
              (throw 'invalid-EOC-length))
            (cons ((asn1-type-aggregator type) (reverse value)) (+ pos 2)))
          (let* ((res (decode-asn1-aux bv pos type))
                 (val (car res))
                 (pos (cdr res)))
            (loop (cons val value) pos)))))

  (let* ((type-byte (bytevector-u8-ref bv pos))
         (constructed? (constructed? type-byte))
         (pos (+ pos 1))
         (len-byte (bytevector-u8-ref bv pos))
         (len (decode-len bv pos))
         (pos (+ pos (lenlen bv pos))))
    (unless (equal? (- type-byte (if constructed? #x20 0)) (asn1-type-type type))
      (throw 'invalid-type-byte type-byte type))
    (when (and constructed? (not (asn1-type-constructed? type)))
      (throw 'invalid-constructed-type))
    (when (and (not constructed?) (not (asn1-type-primitive? type)))
      (throw 'invalid-primitive-type))
    (if (and constructed? (equal? len-byte #x80))
      (decode-asn1-array bv pos type)
      (cons ((asn1-type-parser type) bv pos len) (+ pos len)))))

(define (decode-asn1 bv type)
  "Decode the content of @var{bv}, a DER record of @var{type}"
  (match (decode-asn1-aux bv 0 type)
    ((value . pos) value)))

(define (encode-asn1 type value)
  "Encode the @var{value} of @var{type} into a DER record represented by a
bytevector."
  (let* ((type-byte (asn1-type-type type))
         (type-byte (if (not (asn1-type-primitive? type))
                        (+ #x20 type-byte)
                        type-byte))
         (value ((asn1-type-producer type) value))
         (len (encode-len (bytevector-length value)))
         (type-bv (make-bytevector 1)))
    (bytevector-u8-set! type-bv 0 type-byte)
    (bv-append type-bv len value)))

(define (bv-append . bvs)
  "Append multiple bytevectors in a single bytevector."
  (match bvs
    (() (make-bytevector 0))
    ((bv bvs ...)
     (let* ((bvs (apply bv-append bvs))
            (len1 (bytevector-length bv))
            (len2 (bytevector-length bvs))
            (result (make-bytevector (+ len1 len2))))
       (bytevector-copy! bv 0 result 0 len1)
       (bytevector-copy! bvs 0 result len1 len2)
       result))))

(define (encode-len len)
  "Encode the length in a bytevector, following the ASN.1 specification: on 7
bits if it fits, or the first byte has its most significant bit set, and the 7
remaining bits represent the number of bytes needed to represent the length."
  (define (nbytes len)
    "Return the number of bytes needed to represent a number"
    (if (< len 256) 1 (+ 1 (nbytes (quotient len 256)))))

  (let ((bv (make-bytevector (+ 1 (if (> len 127) (nbytes len) 0)))))
    (bytevector-u8-set! bv 0 (if (> len 127) (+ 128 (nbytes len)) len))
    (when (> len 127)
      (let loop ((val len) (pos (- (bytevector-length bv) 1)))
        (when (> val 0)
          (bytevector-u8-set! bv pos (modulo val 256))
          (loop (quotient val 256) (- pos 1)))))
    bv))