(module (openssl cipher)
  (
   cipher-list
   cipher-by-name
   cipher-key-length
   cipher-iv-length
   cipher-block-size
   cipher-name
   max-key-length
   max-iv-length
   max-block-length
   cipher-context-allocate!
   cipher-context-free!
   cipher-context-reset!
   cipher-context-init!
   cipher-context-update!
   cipher-context-final!
   cipher-context-get-tag
   string-cipher
   string-encrypt-and-digest
   string-decrypt-and-verify
   file-cipher
   open-cipher-port
   )

(import scheme)
(import (chicken base))
(import (chicken bitwise))
(import (chicken blob))
(import (chicken condition))
(import (chicken file posix))
(import (chicken foreign))
(import (chicken format))
(import (chicken gc))
(import (chicken memory))
(import (chicken port))

#>
#include <openssl/err.h>
#include <openssl/evp.h>
<#

(define ERR_clear_error (foreign-lambda void "ERR_clear_error"))
(define ERR_get_error (foreign-lambda unsigned-long "ERR_get_error"))
(define ERR_lib_error_string (foreign-lambda c-string "ERR_lib_error_string" unsigned-long))
(define ERR_func_error_string (foreign-lambda c-string "ERR_func_error_string" unsigned-long))
(define ERR_reason_error_string (foreign-lambda c-string "ERR_reason_error_string" unsigned-long))

(define OBJ_NAME_ALIAS (foreign-value "OBJ_NAME_ALIAS" int))
(define EVP_CIPH_FLAG_AEAD_CIPHER (foreign-value "EVP_CIPH_FLAG_AEAD_CIPHER" int))
(define EVP_CIPH_CCM_MODE (foreign-value "EVP_CIPH_CCM_MODE" int))
(define EVP_CTRL_AEAD_SET_IVLEN (foreign-value "EVP_CTRL_AEAD_SET_IVLEN" int))
(define EVP_CTRL_AEAD_SET_TAG (foreign-value "EVP_CTRL_AEAD_SET_TAG" int))
(define EVP_CTRL_AEAD_GET_TAG (foreign-value "EVP_CTRL_AEAD_GET_TAG" int))
(define EVP_CTRL_CCM_SET_MSGLEN (foreign-value "EVP_CTRL_CCM_SET_MSGLEN" int))

(define-foreign-type OBJ_NAME* (const (c-pointer (struct "obj_name_st"))))
(define-foreign-type EVP_CIPHER* (const c-pointer))
(define-foreign-type EVP_CIPHER_CTX* c-pointer)
(define-foreign-type int* (c-pointer int))

(define evp-ciphers '())
(define-external (EVP_CipherList_callback (OBJ_NAME* obj) (c-pointer _arg)) c-pointer
  (let ((name ((foreign-lambda* c-string ((OBJ_NAME* obj)) "C_return(obj->name);") obj))
        (alias ((foreign-lambda* int ((OBJ_NAME* obj)) "C_return(obj->alias);") obj)))
    (when (not (= alias OBJ_NAME_ALIAS))
      (set! evp-ciphers (cons name evp-ciphers)))
    #f))

(define EVP_CIPHER_CTX_new (foreign-lambda EVP_CIPHER_CTX* "EVP_CIPHER_CTX_new"))
(define EVP_CIPHER_CTX_free (foreign-lambda void "EVP_CIPHER_CTX_free" EVP_CIPHER_CTX*))
(define EVP_CIPHER_CTX_reset (foreign-lambda bool "EVP_CIPHER_CTX_reset" EVP_CIPHER_CTX*))
(define EVP_CIPHER_CTX_set_padding (foreign-lambda bool "EVP_CIPHER_CTX_set_padding" EVP_CIPHER_CTX* bool))
(define EVP_CIPHER_CTX_key_length (foreign-lambda int "EVP_CIPHER_CTX_key_length" EVP_CIPHER_CTX*))
(define EVP_CIPHER_CTX_set_key_length (foreign-lambda bool "EVP_CIPHER_CTX_set_key_length" EVP_CIPHER_CTX* int))
(define EVP_CIPHER_CTX_iv_length (foreign-lambda int "EVP_CIPHER_CTX_iv_length" EVP_CIPHER_CTX*))
(define EVP_CIPHER_CTX_ctrl (foreign-lambda int "EVP_CIPHER_CTX_ctrl" EVP_CIPHER_CTX* int int blob))
(define EVP_CIPHER_CTX_flags (foreign-lambda int "EVP_CIPHER_CTX_flags" EVP_CIPHER_CTX*))
(define EVP_CIPHER_CTX_cipher (foreign-lambda EVP_CIPHER* "EVP_CIPHER_CTX_cipher" EVP_CIPHER_CTX*))

(define EVP_CipherInit_ex (foreign-lambda bool "EVP_CipherInit_ex" EVP_CIPHER_CTX* EVP_CIPHER* c-pointer blob blob int))
(define EVP_CipherUpdate (foreign-lambda bool "EVP_CipherUpdate" EVP_CIPHER_CTX* blob int* (const blob) int))
(define EVP_CipherFinal_ex (foreign-lambda bool "EVP_CipherFinal_ex" EVP_CIPHER_CTX* blob int*))

(define EVP_get_cipherbyname (foreign-lambda EVP_CIPHER* "EVP_get_cipherbyname" (const c-string)))
(define EVP_CIPHER_key_length (foreign-lambda int "EVP_CIPHER_key_length" EVP_CIPHER*))
(define EVP_CIPHER_iv_length (foreign-lambda int "EVP_CIPHER_iv_length" EVP_CIPHER*))
(define EVP_CIPHER_block_size (foreign-lambda int "EVP_CIPHER_block_size" EVP_CIPHER*))
(define EVP_CIPHER_name (foreign-lambda c-string "EVP_CIPHER_name" EVP_CIPHER*))

(define EVP_MAX_KEY_LENGTH (foreign-value "EVP_MAX_KEY_LENGTH" int))
(define EVP_MAX_IV_LENGTH (foreign-value "EVP_MAX_IV_LENGTH" int))
(define EVP_MAX_BLOCK_LENGTH (foreign-value "EVP_MAX_BLOCK_LENGTH" int))

(define (openssl-type-error loc expected #!rest args)
  (abort
   (condition `(exn message ,(format "expected ~a, got" expected)
                    location ,loc
                    arguments ,args)
              '(type))))

(define (openssl-error loc #!rest args)
  (let* ((err (ERR_get_error))
         (message (format "error: library=~a, function=~a, reason=~a"
                          (or (ERR_lib_error_string err) "<unknown>")
                          (or (ERR_func_error_string err) "<unknown>")
                          (or (ERR_reason_error_string err) "<unknown>"))))
    (abort
     (condition `(exn message ,message location ,loc arguments ,args)
                '(i/o)
                `(openssl status #f)))))

(define (cipher-list)
  ;; HACK: without this, the cipher list is empty
  (foreign-code "OPENSSL_init_crypto(OPENSSL_INIT_ADD_ALL_CIPHERS, NULL);")
  (set! evp-ciphers '())
  ((foreign-safe-lambda* void ()
     "OBJ_NAME_do_all_sorted(OBJ_NAME_TYPE_CIPHER_METH, (void(*)(const OBJ_NAME*,void*))EVP_CipherList_callback, NULL);"))
  (reverse evp-ciphers))

(define (cipher-by-name name)
  (EVP_get_cipherbyname name))

(define (cipher-key-length cipher)
  (EVP_CIPHER_key_length cipher))

(define (cipher-iv-length cipher)
  (EVP_CIPHER_iv_length cipher))

(define (cipher-block-size cipher)
  (EVP_CIPHER_block_size cipher))

(define (cipher-name cipher)
  (EVP_CIPHER_name cipher))

(define max-key-length EVP_MAX_KEY_LENGTH)
(define max-iv-length EVP_MAX_IV_LENGTH)
(define max-block-length EVP_MAX_BLOCK_LENGTH)

(define-record cipher-context ptr mode tag-length)

(define (cipher-context-free! context)
  (and-let* ((ctx (cipher-context-ptr context)))
    (EVP_CIPHER_CTX_free ctx)
    (cipher-context-mode-set! context #f)
    (cipher-context-ptr-set! context #f)))

(define (cipher-context-allocate!)
  (ERR_clear_error)
  (let ((ctx (EVP_CIPHER_CTX_new)))
    (when (not ctx)
      (openssl-error 'cipher-context-allocate!))
    (set-finalizer! (make-cipher-context ctx #f #f) cipher-context-free!)))

(define (cipher-context-unwrap! context)
  (let ((ctx (cipher-context-ptr context)))
    (when (not ctx)
      (openssl-type-error 'cipher-context-unwrap! "valid context pointer" #f))
    ctx))

(define (cipher-context-reset! context)
  (let ((ctx (cipher-context-unwrap! context)))
    (ERR_clear_error)
    (when (not (EVP_CIPHER_CTX_reset ctx))
      (openssl-error 'cipher-context-reset!))
    (cipher-context-tag-length-set! context #f)
    (void)))

(define (aead-cipher? ctx)
  (bitwise-and (EVP_CIPHER_CTX_flags ctx) EVP_CIPH_FLAG_AEAD_CIPHER))

(define (remove-known-key-args known-keys args)
  (let loop ((args args)
             (unknown-args '()))
    (if (null? args)
        (reverse unknown-args)
        (let ((arg (car args)))
          (if (memv arg known-keys)
              (if (pair? (cdr args))
                  (loop (cddr args) unknown-args)
                  (loop (cdr args) unknown-args))
              (loop (cdr args) (cons arg unknown-args)))))))

(define (check-unknown-key-args known-keys args)
  (let ((unknown-args (remove-known-key-args known-keys args)))
    (when (pair? unknown-args)
      (fprintf (current-error-port) "warning: unknown rest arguments: ~s\n" unknown-args))))

(define (cipher-context-init! context cipher key iv
                              #!rest args
                              #!key
                              (mode 'encrypt)
                              (padding #t)
                              (effective-key-length #f)
                              (auth-data #f)
                              (tag-length #f)
                              (expected-tag #f)
                              (effective-iv-length #f)
                              (message-length #f))
  (define (mode->flag mode)
    (case mode
      ((encrypt) 1)
      ((decrypt) 0)
      (else (openssl-type-error 'cipher-context-init! "mode symbol" (list mode)))))
  (define (ccm-cipher? ctx)
    (bitwise-and (EVP_CIPHER_CTX_flags ctx) EVP_CIPH_CCM_MODE))
  (check-unknown-key-args '(#:mode #:padding #:effective-key-length #:auth-data
                            #:tag-length #:expected-tag #:effective-iv-length
                            #:message-length) args)
  (let ((ctx (cipher-context-unwrap! context))
        (key-length (or effective-key-length (blob-size key)))
        (iv-length (and iv (blob-size iv))))
    (ERR_clear_error)
    (when (not (EVP_CipherInit_ex ctx cipher #f #f #f (mode->flag mode)))
      (openssl-error 'cipher-context-init! (list cipher mode)))
    (cipher-context-mode-set! context mode)
    (when (> key-length (blob-size key))
      (openssl-type-error "effective key length <= key size" key-length (blob-size key)))
    (when (not (EVP_CIPHER_CTX_set_key_length ctx key-length))
      (openssl-error 'cipher-context-init! (list key-length effective-key-length)))
    (let ((min-iv-length (or effective-iv-length (EVP_CIPHER_CTX_iv_length ctx))))
      (when (and iv-length (< iv-length min-iv-length))
        (openssl-type-error 'cipher-context-init! "sufficient iv length" iv-length min-iv-length)))
    (when effective-iv-length
      (when (not (aead-cipher? ctx))
        (openssl-type-error 'cipher-context-init! "AEAD cipher" (cipher-name cipher)))
      (when (> effective-iv-length max-iv-length)
        (openssl-type-error 'cipher-context-init! "integer <= 16" effective-iv-length))
      (when (not (EVP_CIPHER_CTX_ctrl ctx EVP_CTRL_AEAD_SET_IVLEN effective-iv-length #f))
        (openssl-error 'cipher-context-init! (list effective-iv-length))))
    (when (and expected-tag tag-length)
      (when (not (aead-cipher? ctx))
        (openssl-type-error 'cipher-context-init! "AEAD cipher" (cipher-name cipher)))
      (when (not (eqv? mode 'decrypt))
        (openssl-type-error 'cipher-context-init! "decrypt mode" mode))
      (when (not tag-length)
        (openssl-type-error 'cipher-context-init! "tag length" #f))
      (when (> tag-length (blob-size expected-tag))
        (openssl-type-error 'cipher-context-init! "tag shorter than tag length" tag-length (blob-size expected-tag)))
      (when (not (EVP_CIPHER_CTX_ctrl ctx EVP_CTRL_AEAD_SET_TAG tag-length expected-tag))
        (openssl-error 'cipher-context-init! (list expected-tag tag-length))))
    (when (and (not expected-tag) tag-length)
      (when (not (aead-cipher? ctx))
        (openssl-type-error 'cipher-context-init! "AEAD cipher" (cipher-name cipher)))
      (when (not (eqv? mode 'encrypt))
        (openssl-type-error 'cipher-context-init! "encrypt mode" mode))
      (when (> tag-length max-iv-length)
        (openssl-type-error 'cipher-context-init! "integer <= 16" tag-length))
      (when (not (EVP_CIPHER_CTX_ctrl ctx EVP_CTRL_AEAD_SET_TAG tag-length #f))
        (openssl-error 'cipher-context-init! (list tag-length)))
      (cipher-context-tag-length-set! context tag-length))
    (when (not (EVP_CipherInit_ex ctx #f #f key iv -1))
      (openssl-error 'cipher-context-init! (list cipher key iv)))
    (when message-length
      (when (not (ccm-cipher? ctx))
        (openssl-type-error 'cipher-context-init! "CCM cipher mode" (cipher-name cipher)))
      ;; https://github.com/pyca/cryptography/blob/0034926f2cca02258f50e9faccb90ec344790159/src/cryptography/hazmat/backends/openssl/aead.py#L108
      ;; https://github.com/pyca/cryptography/blob/0034926f2cca02258f50e9faccb90ec344790159/src/cryptography/hazmat/backends/openssl/aead.py#L77
      (let-location ((_length int))
        (when (not (EVP_CipherUpdate ctx #f (location _length) #f message-length))
          (openssl-error 'cipher-context-init! (list message-length)))))
    (when auth-data
      (when (not (aead-cipher? ctx))
        (openssl-type-error 'cipher-context-init! "AEAD cipher" (cipher-name cipher)))
      (let-location ((_length int))
        (when (not (EVP_CipherUpdate ctx #f (location _length) auth-data (blob-size auth-data)))
          (openssl-error 'cipher-context-init! (list auth-data (blob-size auth-data))))))
    (EVP_CIPHER_CTX_set_padding ctx padding)
    (void)))

(define (cipher-context-update! context blob)
  (let ((ctx (cipher-context-unwrap! context))
        (buf (make-blob (+ (blob-size blob) max-block-length))))
    (ERR_clear_error)
    (let-location ((buf-length int))
      (when (not (EVP_CipherUpdate ctx buf (location buf-length) blob (blob-size blob)))
        (openssl-error 'cipher-context-update! (list blob (blob-size blob))))
      (let ((ret (make-blob buf-length)))
        (move-memory! buf ret buf-length)
        ret))))

(define (cipher-context-final! context)
  (let ((ctx (cipher-context-unwrap! context))
        (buf (make-blob max-block-length)))
    (ERR_clear_error)
    (let-location ((buf-length int))
      (when (not (EVP_CipherFinal_ex ctx buf (location buf-length)))
        (openssl-error 'cipher-context-final!))
      (let ((ret (make-blob buf-length)))
        (move-memory! buf ret buf-length)
        ret))))

(define (cipher-context-get-tag context)
  (let ((ctx (cipher-context-unwrap! context)))
    (ERR_clear_error)
    (when (not (aead-cipher? ctx))
      (openssl-type-error 'cipher-context-get-tag "AEAD cipher" (cipher-name (EVP_CIPHER_CTX_cipher ctx))))
    (let ((mode (cipher-context-mode context))
          (tag-length (cipher-context-tag-length context)))
      (when (not mode)
        (openssl-type-error 'cipher-context-get-tag "initialized context" #f))
      (when (not (eqv? mode 'encrypt))
        (openssl-type-error 'cipher-context-get-tag "encrypt mode" mode))
      (when (not tag-length)
        (openssl-type-error 'cipher-context-get-tag "tag length to be set" #f))
      (let ((buf (make-blob tag-length)))
        (when (not (EVP_CIPHER_CTX_ctrl ctx EVP_CTRL_AEAD_GET_TAG tag-length buf))
          (openssl-error 'cipher-context-get-tag tag-length))
        buf))))

(define (string-cipher cipher str key iv #!rest options)
  (let ((context (cipher-context-allocate!)))
    (apply cipher-context-init! context cipher key iv options)
    (let* ((output (cipher-context-update! context (string->blob str)))
           (final (cipher-context-final! context))
           (output (string-append (blob->string output) (blob->string final))))
      (cipher-context-free! context)
      output)))

(define (string-encrypt-and-digest cipher str key iv #!rest options)
  (let ((context (cipher-context-allocate!)))
    (apply cipher-context-init! context cipher key iv mode: 'encrypt options)
    (let* ((output (cipher-context-update! context (string->blob str)))
           (final (cipher-context-final! context))
           (output (string-append (blob->string output) (blob->string final)))
           (tag (blob->string (cipher-context-get-tag context))))
      (cipher-context-free! context)
      (values output tag))))

(define (string-decrypt-and-verify cipher str tag key iv #!rest options)
  (let ((context (cipher-context-allocate!)))
    (apply cipher-context-init! context cipher key iv
           mode: 'decrypt
           tag-length: (string-length tag)
           expected-tag: (string->blob tag)
           options)
    (let* ((output (cipher-context-update! context (string->blob str)))
           (final (cipher-context-final! context))
           (output (string-append (blob->string output) (blob->string final))))
      (cipher-context-free! context)
      output)))

(define (file-cipher cipher in-path out-path key iv #!rest options)
  (let* ((buf-size 4096)
         (buf (make-blob buf-size))
         (context (cipher-context-allocate!))
         (in (file-open in-path open/rdonly))
         (out (file-open out-path
                         (+ open/wronly open/creat open/trunc)
                         (+ perm/irusr perm/iwusr))))
    (apply cipher-context-init! context cipher key iv options)
    (let loop ()
      (let ((count (cadr (file-read in buf-size buf))))
        (when (positive? count)
          (file-write out (cipher-context-update! context buf))
          (loop))))
    (file-write out (cipher-context-final! context))
    (file-close in)
    (file-close out)
    (cipher-context-free! context)))

(define (open-cipher-port cipher out key iv #!rest options)
  (let ((context (cipher-context-allocate!)))
    (apply cipher-context-init! context cipher key iv options)
    (make-output-port
     (lambda (str)
       (display (blob->string (cipher-context-update! context (string->blob str))) out))
     (lambda ()
       (display (blob->string (cipher-context-final! context)) out))
     (lambda ()
       (flush-output out)))))

)