#lang racket/base

(require ffi/unsafe
         ffi/unsafe/define
         racket/match
         openssl
         openssl/libssl
         openssl/libcrypto
         (prefix-in c: racket/contract)
         racket/string
         )

(provide generate-self-signed-cert
         self-signed-cert
         self-signed-cert?
         private-key
         certificate
         x509-cert
         )

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Exported struct
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define-struct self-signed-cert
  (private-key certificate))

(c:define/contract (private-key ssc)
  (c:-> self-signed-cert? string?)
  (self-signed-cert-private-key ssc))

(c:define/contract (certificate ssc)
  (c:-> self-signed-cert? string?)
  (self-signed-cert-certificate ssc))

(define x509-cert certificate)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Supportive macros / functions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define-syntax version-ffi-define
  (syntax-rules (openssl-major libopenssl libssl)
    ((_ version 
        (name definition) ...)
     (begin
       (define name #f)
       ...
       (cond
         ([= openssl-major version]
          (with-handlers
              ([exn:fail?
                (λ (exn) (set! name (get-ffi-obj
                                  (symbol->string 'name) libssl definition)))])
            (set! name (get-ffi-obj (symbol->string 'name) libopenssl definition)))
          ...))
       )
     )
    )
  )

(define-syntax version-define
  (syntax-rules (openssl-major)
    ((_ version
        (name definition)
        ...)
     (begin
       (define name #f)
       ...
       (cond
         ([= openssl-major version]
          (set! name definition)
          ...
          ))
       )
     )
    )
  )

(define (is-ip? h)
  (if (string? h)
      (let ((re #px"^[0-9]+[.][0-9]+[.][0-9]+[.][0-9]+$"))
        (not (eq? (regexp-match re (string-trim h)) #f)))
      #f))
  

(define (is-dns? h)
  (if (string? h)
      (let ((re #px"[^ .]+([.][^ .]+)*"))
        (not (eq? (regexp-match re (string-trim h)) #f)))
      #f))

(define (make-alt-entry host)
  (if (is-ip? host)
      (format "IP:~a" host)
      (format "DNS:~a" host)))

(define (list-of-hosts? h)
  (letrec ((f (λ (l)
                (if (null? l)
                    #t
                    (and (or (is-ip? (car l)) (is-dns? (car l)))
                         (f (cdr l)))))))
    (if (list? h)
        (if (null? h)
            #f
            (f h))
        #f)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; FFI Stuff needed
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define os (system-type 'os*))

(define libopenssl libcrypto)

(define-ffi-definer define-ssl libopenssl)

(define _EVP_PKEY-pointer (_cpointer/null 'EVP_PKEY*))
(define _RSA-pointer (_cpointer/null 'RSA*))
(define _X509-pointer (_cpointer/null 'X509*))
(define _ASN1_INTEGER-pointer (_cpointer/null 'ASN1_INTEGER*))
(define _ASN1_STRING-pointer (_cpointer/null 'ASN1_STRING*))
(define _ASN1_TIME-pointer (_cpointer/null 'ASN1_TIME*))
(define _X509_NAME-pointer (_cpointer/null 'X509_NAME*))
(define _EVP_MD-pointer (_cpointer/null 'EVP_MD*))
(define _BIO_METHOD-pointer (_cpointer/null 'BIO_METHOD*))
(define _BIO-pointer (_cpointer/null 'BIO*))
(define _EVP_CIPHER-pointer (_cpointer/null 'EVP_CYPHER*))
(define _X509_EXTENSION-pointer (_cpointer/null 'X509_EXTENSION*))


(define RSA_F4 #x10001)
(define MBSTRING_FLAG #x1000)
(define MBSTRING_UTF8 MBSTRING_FLAG)
(define MBSTRING_ASC (+ MBSTRING_FLAG 1))
(define BIO_CTRL_INFO 3)
(define NID_rsaEncryption 6)
(define EVP_PKEY_RSA NID_rsaEncryption)
(define V_ASN1_OCTET_STRING   4)
(define NID_subject_alt_name   85)

(define _string/utf-8-pointer (_ptr o _string/utf-8))

;typedef int pem_password_cb(char *buf, int size, int rwflag, void *userdata)
(define _pem_password_cb
  (_fun _string/utf-8 _int _int _pointer -> _int))

(define _gen_rsa_cb
  (_fun _int _int _pointer -> _void))


;; Check which openssl version we're dealing with
(define openssl-major #f)

(with-handlers ([exn:fail? (λ (exn) (set! openssl-major 1))])
  (define-ssl EVP_RSA_gen
    (_fun _int -> _EVP_PKEY-pointer))
  (set! openssl-major 3))


(version-ffi-define 1
  (EVP_PKEY_new (_fun -> _EVP_PKEY-pointer))
  (EVP_PKEY_free (_fun _EVP_PKEY-pointer -> _void))
  (RSA_generate_key (_fun _int _int _gen_rsa_cb _pointer -> _RSA-pointer))
  (EVP_PKEY_assign (_fun _EVP_PKEY-pointer _int _RSA-pointer -> _int))
  (EVP_sha1 (_fun -> _EVP_MD-pointer))
  
  (X509_new (_fun -> _X509-pointer))
  (X509_free (_fun _X509-pointer -> _void))
  (X509_get_serialNumber (_fun _X509-pointer -> _ASN1_INTEGER-pointer))
  (X509_get0_notBefore (_fun _X509-pointer -> _ASN1_TIME-pointer))
  (X509_get0_notAfter (_fun _X509-pointer -> _ASN1_TIME-pointer))
  (X509_gmtime_adj (_fun _ASN1_TIME-pointer _long -> _ASN1_TIME-pointer))
  (X509_set_pubkey (_fun _X509-pointer _EVP_PKEY-pointer -> _int))
  (X509_get_subject_name (_fun _X509-pointer -> _X509_NAME-pointer))
  (X509_NAME_add_entry_by_txt (_fun _X509_NAME-pointer _string/utf-8 _int _string/utf-8 _int _int _int -> _int))
  (X509_set_issuer_name (_fun _X509-pointer _X509_NAME-pointer -> _int))
  (X509_sign (_fun _X509-pointer _EVP_PKEY-pointer _EVP_MD-pointer -> _int))
  (X509_EXTENSION_create_by_NID (_fun _pointer ; could also be, (ep : (_ptr o _X509_EXTENSION-pointer)), but works fine when #f is provided
                                      _int _int _ASN1_STRING-pointer -> (p : _X509_EXTENSION-pointer)
                                      -> p))
  (X509_add_ext (_fun _X509-pointer _X509_EXTENSION-pointer _int -> _int))
  (X509_EXTENSION_free (_fun _X509_EXTENSION-pointer -> _void))
  
  
  (ASN1_INTEGER_set (_fun _ASN1_INTEGER-pointer _long -> _int))
  (ASN1_STRING_new (_fun -> _ASN1_STRING-pointer))
  (ASN1_STRING_free (_fun _ASN1_STRING-pointer -> _void))
  (ASN1_STRING_type_new (_fun _int -> _ASN1_STRING-pointer))
  (ASN1_OCTET_STRING_set (_fun _ASN1_STRING-pointer _string/utf-8 _int -> _int))

  (BIO_s_mem (_fun -> _BIO_METHOD-pointer))
  (BIO_ctrl (_fun _BIO-pointer _int _long
                  (out : (_ptr o _bytes)) -> (len : _long) -> (list len out)))

  (PEM_write_bio_PrivateKey (_fun _BIO-pointer _EVP_PKEY-pointer _EVP_CIPHER-pointer
                                  _string/utf-8 _int _pem_password_cb _pointer -> _int))
  (PEM_write_bio_X509 (_fun _BIO-pointer _X509-pointer -> _int))
  (BIO_new (_fun _BIO_METHOD-pointer -> _BIO-pointer))
  (BIO_puts (_fun _BIO-pointer _string/utf-8 -> _int))
  (BIO_free (_fun _BIO-pointer -> _int))
  )

(version-define 1
   (BIO_get_mem_data (λ (bio-ptr)
                       (let ((r (BIO_ctrl bio-ptr BIO_CTRL_INFO 0)))
                         (cadr r))))
   (ASN1_OCTET_STRING_new (λ ()
                            (ASN1_STRING_type_new V_ASN1_OCTET_STRING)))
   )

(cond
  ((= openssl-major 3)
   (error "OpenSSL major version 3 is not supported yet")))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Provided function
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(version-define 1
   (generate-key
    (λ (bits)
      (let ((pkey (EVP_PKEY_new)))
        (when (eq? pkey #f)
          (error "Unable to create EVP_PKEY structure"))

        (let* ((rsa (RSA_generate_key bits RSA_F4 #f #f)))
          (when (= (EVP_PKEY_assign pkey EVP_PKEY_RSA rsa) 0)
            (EVP_PKEY_free pkey)
            (error "Unable to generate RSA key"))

          pkey))))

   (generate-x509
    (λ (pkey duration-in-days country company hosts)
      (let ((x509 (X509_new)))
        (when (eq? x509 #f)
            (error "Unable to create X509 structure"))

        (ASN1_INTEGER_set (X509_get_serialNumber x509) 1)
        (X509_gmtime_adj (X509_get0_notBefore x509) 0)
        (X509_gmtime_adj (X509_get0_notAfter x509) (* duration-in-days 24 3600))
        (X509_set_pubkey x509 pkey)

        (let* ((x509-name (X509_get_subject_name x509))
               (first-host (car hosts)))
          (X509_NAME_add_entry_by_txt x509-name
                                      "C" MBSTRING_UTF8 country -1 -1 0)
          (X509_NAME_add_entry_by_txt x509-name
                                      "O" MBSTRING_UTF8 company -1 -1 0)
          (X509_NAME_add_entry_by_txt x509-name
                                      "CN" MBSTRING_UTF8 first-host -1 -1 0)
          (X509_set_issuer_name x509 x509-name)

          (let* ((alt-name (string-join
                           (map make-alt-entry hosts) ", "))
                 (ext-san #f)
                 (subj-alt-name-asn1 #f)
                 )

            (set! subj-alt-name-asn1 (ASN1_OCTET_STRING_new))
            (when (eq? subj-alt-name-asn1 #f)
              (error "Cannot allocate Subject Alt Name ASN1 string"))

            (ASN1_OCTET_STRING_set subj-alt-name-asn1
                                   alt-name (string-length alt-name))

            (let ((r (X509_EXTENSION_create_by_NID #f NID_subject_alt_name 0 subj-alt-name-asn1)))
              (when (eq? r #f)
                (error "Cannot allocate X509 Extenstion for Subject Alt Name"))

              (let* ((extension_san r)
                     (re (X509_add_ext x509 extension_san -1)))
                (when (= re 0)
                    (error "Cannot add extension to X509"))

                (X509_EXTENSION_free extension_san)))

            (ASN1_STRING_free subj-alt-name-asn1)
            )
          )

        (when (= (X509_sign x509 pkey (EVP_sha1)) 0)
          (X509_free x509)
          (error "Error signing certificate"))

        x509)))
          
   (generate-self-signed-cert*
    (λ (bits duration-in-days hosts country company)
      (let* ((pkey (generate-key bits))
             (x509 (generate-x509 pkey duration-in-days country company hosts))
             (pkey-data #f)
             (x509-data #f)
             )
        
        (let ((bio (BIO_new (BIO_s_mem))))
          (let ((r (PEM_write_bio_PrivateKey bio pkey #f #f 0 #f #f)))
            (when (= r 0)
              (BIO_free bio)
              (error "Unable to write private key to memory"))
            (let ((data (BIO_get_mem_data bio)))
              (set! pkey-data data))
            (BIO_free bio)))

        (let ((bio (BIO_new (BIO_s_mem))))
          (let ((r (PEM_write_bio_X509 bio x509)))
            (when (= r 0)
              (BIO_free bio)
              (error "Unable to write X.509 certificate to memory"))
            (let ((data (BIO_get_mem_data bio)))
              (set! x509-data data))
            (BIO_free bio)))

        (EVP_PKEY_free pkey)
        (X509_free x509)

        ;(displayln pkey-data)
        ;(displayln x509-data)
        ;(displayln (format "pkey: ~a" (bytes->string/utf-8 pkey-data)))
        ;(displayln (format "cert: ~a" (bytes->string/utf-8 x509-data)))

        (make-self-signed-cert pkey-data x509-data)
        ;(bytes->string/utf-8 pkey-data)
        ;                       (bytes->string/utf-8 x509-data))
        )
      )
    )
   )

(c:define/contract (generate-self-signed-cert bits duration-in-days hosts country company)
  (c:-> integer? integer? (c:or/c is-ip? is-dns? list-of-hosts?) string? string?
        self-signed-cert?)
  (if (eq? generate-self-signed-cert* #f)
      (error "No openssl FFI glue code available")
      (let ((h (if (list-of-hosts? hosts) hosts (list hosts))))
        (generate-self-signed-cert* bits duration-in-days h country company)
        )
      )
  )