(module crypto-tools
(blob->hexstring blob->hexstring/uppercase hexstring->blob
 blob-xor blob-pad blob-unpad blob-pkcs5-pad blob-pkcs5-unpad no-pad no-unpad
 make-cbc-encryptor make-cbc-decryptor make-cbc*-encryptor make-cbc*-decryptor
 make-ctr-encryptor make-ctr-decryptor make-ctr*-encryptor make-ctr*-decryptor)

(import scheme)
(import (chicken foreign))
(import (chicken blob))
(import (chicken base))
(import (chicken format))
(import (chicken locative))
(import (chicken memory))

;zero sizes blob
(define *the-null-blob* (make-blob 0))

;ensures the input is a blob of certain size
(define (check-blob blob len function)
   (if (blob? blob)
      (if (= (blob-size blob) len)
         (void)
         (error (sprintf "~A: Input blob was ~A bytes long, when ~A were required" function (blob-size blob) len)))
      (error (sprintf "~A: Input was not a blob" function))))

;allocates a new string where every byte of the blob is encoded as a two char hex number
(define (blob->hexstring blob)
   (if (zero? (blob-size blob))
      ""
      (let ((str (make-string (* 2 (blob-size blob)))))
         ((foreign-lambda* void ((blob in) (int in_len) ((c-pointer char) out)) "
            while (in_len--) {
               *out++ = \"0123456789abcdef\"[(*in) >> 4];
               *out++ = \"0123456789abcdef\"[(*in) & 0x0f];
               in++;
            }
            ") blob (blob-size blob) (make-locative str))
         str)))

;same as blob->hexstring but with uppercase hex chars
(define (blob->hexstring/uppercase blob)
   (if (zero? (blob-size blob))
      ""
      (let ((str (make-string (* 2 (blob-size blob)))))
         ((foreign-lambda* void ((blob in) (int in_len) ((c-pointer char) out)) "
            while (in_len--) {
               *out++ = \"0123456789ABCDEF\"[(*in) >> 4];
               *out++ = \"0123456789ABCDEF\"[(*in) & 0x0f];
               in++;
            }
            ") blob (blob-size blob) (make-locative str))
         str)))

;converts a hex encoded string of two chars per byte to a blob
(define (hexstring->blob string)
   (let ((len (string-length string)))
      (cond
         ((zero? len) *the-null-blob*)
         ((odd? len)
            (error "hexstring->blob: The supplied string must have an even length"))
         (else
            (let ((blob (make-blob (/ (string-length string) 2))))
               (if ((foreign-lambda* bool (((c-pointer char) in) (int in_len) (blob out)) "
                  while (in_len) {
                     if (*in >= '0' && *in <= '9') *out = (*in - '0') << 4;
                     else if (*in >= 'A' && *in <= 'F') *out = (*in - 'A' + 10) << 4;
                     else if (*in >= 'a' && *in <= 'f') *out = (*in - 'a' + 10) << 4;
                     else C_return(1);
                     in++;
            
                     if (*in >= '0' && *in <= '9') *out |= (*in - '0');
                     else if (*in >= 'A' && *in <= 'F') *out |= (*in - 'A' + 10);
                     else if (*in >= 'a' && *in <= 'f') *out |= (*in - 'a' + 10);
                     else C_return(1);
                     in++;
            
                     out++;
            
                     in_len -= 2;
                  }
                  C_return(0);
                  ") (make-locative string) (string-length string) blob)
                  (error "hexstring->blob: Invalid character(s) in input string")
                  blob))))))

;makes a new blob and copies in it part of the contents of the original
(define (subblob blob offset length)
   (if (not (blob? blob)) (error "subblob: Argument must be a blob"))
   (if (> (+ offset length) (blob-size blob))
       (error (sprintf "subblob: Argument must be large enough: ~A bytes supplied, but cutting out ~A bytes from ~A would need ~A" (blob-size blob) length offset (+ length offset))))
       (cond
         ((and (zero? offset) (= (blob-size blob) length))
          blob)
         ((zero? length)
          *the-null-blob*)
         (else
           (let ((out (make-blob length)))
             (move-memory! blob out length offset 0)
              out))))

;makes a new blob with contetns the xor of the arguments up to the length of the smallest one.
(define (blob-xor a b)
   (if (not (blob? a)) (error "blob-xor: Arguments must be blobs"))
   (if (= (blob-size a) 0) b)
   (if (= (blob-size b) 0) a)
   ;(check-blob b (blob-size a) 'blob-xor)
   (let* ((minlen (min (blob-size a) (blob-size b)))
          (out (make-blob minlen)))
      ((foreign-lambda* void ((int len) (blob a) (blob b) (blob out)) "
         while (len--) {
            (*out++) = (*a++) ^ (*b++);
         }
      ") minlen a b out)
      out))

;pads a blob with the a value of the number of bytes of that pad.
;in case a whole blocklength was added then the padvalue is that blocklength.
;that means that a pad is always present.
(define (blob-pkcs5-pad in len)
  (if (not (blob? in)) (error "blob-pkcs5-pad: Argument must be a blob"))
  (let* ((inlen (blob-size in))
	 (extrabytes (remainder inlen len))
	 (padsz (if (zero? extrabytes) len extrabytes))
	 (totlen (+ inlen padsz))
	 (out (make-blob (+ inlen padsz))))
    (move-memory! in out inlen 0 0)
    ((foreign-lambda* void ((int inlen) (int outlen) (int padnum) (blob out)) "
       while (inlen < outlen) {
           out[inlen++ - 1] = (char)padnum;
       }
    ") inlen totlen padsz out)
    out))

;unpads by reading the last byte and removing as many as its value
(define (blob-pkcs5-unpad in)
  (if (not (blob? in)) (error "blob-pkcs5-unpad: Argument must be a blob"))
  (let* ((inlen (blob-size in))
         (padsz ((foreign-lambda* int ((blob in) (int inlen)) "
		   char sz;
		   sz = in[inlen-1];
                   C_return((int)sz);
                ") in inlen)))
    (if (or (<= padsz 0) (> padsz inlen)) (error "blob-pkcs5-unpad: Argument must be a validly padded blob. Got:" padsz " insize:" inlen))
    (let* ((outlen (- inlen padsz))
           (out (make-blob outlen)))
      (move-memory! in out outlen)
      out)))

;noop. returns the input as is without padding
(define (no-pad in len) in)
;noop, returns the input without unpadding
(define (no-unpad in) in)

;pads by adding the byte 0x80 followed by 0x00 up to the desired length.
(define (blob-pad in len)
   (if (not (blob? in)) (error "blob-pad: Argument must be a blob"))
   (if (>= (blob-size in) len) (error "blob-pad: Argument must be smaller than the block size"))
   (let ((out (make-blob len))
         (inlen (blob-size in)))
      (move-memory! in out inlen 0 0)
      ((foreign-lambda* void ((int inlen) (int outlen) (blob in) (blob out)) "
         out[inlen++] = 0x80;
         while (inlen < outlen) {
            out[inlen++] = 0x00;
         }
      ") inlen len in out)
      out))

;unpads by removing 0x00 until it discovers a 0x80 byte
(define (blob-unpad in)
   (if (not (blob? in)) (error "blob-unpad: Argument must be a blob"))
   (if (< (blob-size in) 1) (error "blob-unpad: Argument must be at least a byte long"))
   (let* ((inlen (blob-size in))
          (outlen
             ((foreign-lambda* int ((int inlen) (blob in)) "
                while (inlen > 0 && in[inlen] == 0x00) {
                  inlen--;
                }
                if (in[inlen] != 0x80) C_return(-1);
                C_return(inlen);
            ") (- inlen 1) in)))
         
         (if (= outlen -1) (error "blob-unpad: Argument must be a validly padded blob"))
         (if (zero? outlen) *the-null-blob*)
         (let ((out (make-blob outlen)))
            (move-memory! in out outlen)
            out)))

;holds the state of the counter needed for ctr-mode. On every invocation of the resulting
;function the value of the counter will be increased by one.On completion it should wrap around.
(define (make-next-ctr iv bsize)
          (let ((cnt iv)
                (incblob (foreign-lambda* void ((blob in) (blob out) (int blocksize)) "
                           unsigned char cur;
                           int block;
                           memmove(out,in,blocksize);
                           for(block=blocksize-1;block>=0;block--) {
                                cur = in[block];
                                if (cur >= 0xFF) {
                                      out[block] = 0;
                                      continue;
                                }
                                out[block] = cur+1;
                                break;
                         }
                     ")))
            (lambda ()
                (let* ((ocnt cnt)
                      (ncnt (make-blob bsize)))
                      (incblob ocnt ncnt bsize)
                      (set! cnt ncnt)
                      ocnt))))

;internal ctr-encrypt function used by the make- variants
;on each block fetches the current counter state (which increases the counter)
;encrypts it and xor as much blob as it has available.
(define (ctr-encrypt input inoffset inputsize nextfn output outoffset encryptor blocksize) 
  (cond
     ((= inoffset inputsize) output) ; we just output what we have.
     ((<= (+ inoffset blocksize) inputsize) ; Just another block
        (let* ((inblock (subblob input inoffset blocksize))
               (outblock (blob-xor inblock (encryptor (nextfn)))))
           (move-memory! outblock output blocksize 0 outoffset)
           (ctr-encrypt input (+ inoffset blocksize) inputsize nextfn output (+ outoffset blocksize) encryptor blocksize))) ; Recurse
     (else ; Partial block
        (let* ((inblock (subblob input inoffset (- inputsize inoffset)))
               (outblock (blob-xor inblock (encryptor (nextfn)))))
           (move-memory! outblock output (blob-size outblock) 0 outoffset)
           output))))

;creates a ctr encryptor that needs an iv and input blob 
(define (make-ctr-encryptor encryptor blocksize)
  (lambda (input iv)
     (let* ((inputsize (blob-size input))
            (outputsize inputsize) ;ct size since no iv and no padding.
            (output (make-blob outputsize))
            (nextf (make-next-ctr iv blocksize)))
        (ctr-encrypt input 0 inputsize nextf output 0 encryptor blocksize))))

;creates a ctr encryptor function that needs an iv and input blob but
;stores the IV in the first output block (unencrypted since ctr iv is only a counter)
(define (make-ctr*-encryptor encryptor blocksize)
  (lambda (input iv)
     (let* ((inputsize (blob-size input))
            (outputsize (+ inputsize blocksize)) ;iv + ct cause ctr doesn't need padding.
            (output (make-blob outputsize))
            (nextf (make-next-ctr iv blocksize)))
        (move-memory! iv output blocksize)
        (ctr-encrypt input 0 inputsize nextf output blocksize encryptor blocksize))))

;internal ctr-decrypt function used by the make- variants
;gets the current state of the counter (which increases the counter) and then 
;*encrypts* the counter and xors that with as much blob as it has.
(define (ctr-decrypt input inoffset inputsize nextfn output outoffset encryptor blocksize)
  (if (> (+ inoffset blocksize) inputsize)
     ; Last block
     (let* ((inblock (subblob input inoffset (- (blob-size input) inoffset)))
            (outblock
                (blob-xor inblock
                  (encryptor (nextfn)))))
           (move-memory! outblock output (blob-size outblock) 0 outoffset)
           (subblob output 0 (+ outoffset (- (blob-size input) inoffset)))) ; Terminate
     ; More blocks following
     (let* ((inblock (subblob input inoffset blocksize))
            (outblock
               (blob-xor inblock
                 (encryptor (nextfn)))))
           (move-memory! outblock output blocksize 0 outoffset)
           (ctr-decrypt input (+ inoffset blocksize) inputsize nextfn output (+ outoffset blocksize) encryptor blocksize)))) ; Recurse

;returns a function that requires an input and an iv. it starts decrypting from the first blob.
;typically ctr mode decryption requires an *encryptor*
(define (make-ctr-decryptor encryptor blocksize)
  (lambda (input iv)
     (let* ((inputsize (blob-size input))
            (output (make-blob inputsize))
            (nextf (make-next-ctr iv blocksize)))
        (ctr-decrypt input 0 inputsize nextf output 0 encryptor blocksize))))

;returns a function that requires an input. it reads the unencrypted iv from the first block
;and then uses that to start decrypting in ctr mode the blob that follows.
;typically ctr mode decryption requires an *encryptor*
(define (make-ctr*-decryptor encryptor blocksize)
  (lambda (input)
     (let* ((inputsize (blob-size input))
            (output (make-blob inputsize))
            (iv (subblob input 0 blocksize))
            (nextf (make-next-ctr iv blocksize)))
        (ctr-decrypt input blocksize inputsize nextf output 0 encryptor blocksize))))

;internal cbc encrypt function used by the make- variants. cbc tradionally requires padding
(define (cbc-encrypt input inoffset inputsize iv output outoffset padfn encryptor blocksize)
    (cond
      ((= inoffset inputsize) ; Zero bytes
         (let* ((inblock (padfn *the-null-blob* blocksize))
                (outblock (encryptor (blob-xor iv inblock))))
            (move-memory! outblock output blocksize 0 outoffset)
            output)) ; Terminate
      ((<= (+ inoffset blocksize) inputsize) ; Just another block
         (let* ((inblock (subblob input inoffset blocksize))
                (outblock (encryptor (blob-xor iv inblock))))
            (move-memory! outblock output blocksize 0 outoffset)
            (cbc-encrypt input (+ inoffset blocksize) inputsize outblock output (+ outoffset blocksize) padfn encryptor blocksize))) ; Recurse
      (else ; Partial block
         (let* ((inblock (padfn
                  (subblob input inoffset (- inputsize inoffset))
                  blocksize))
                (outblock (encryptor (blob-xor iv inblock))))
            (move-memory! outblock output blocksize 0 outoffset)
            output)))) ; Terminate

;creates a cbs encryptor function that requires an input and an iv in order to start the
;cbc decryption from the first blob. In the make arguments a padding function like pkcs5 or pad-blob is needed.
(define (make-cbc-encryptor encryptor padfn blocksize)
  (lambda (input iv)
    (let* ((inputsize (blob-size input))
           (output-whole-blocks (quotient inputsize blocksize))
           (output-overflow (remainder inputsize blocksize))
           (outputsize (if (zero? output-overflow)
              (+ inputsize blocksize)
              (* (+ 1 output-whole-blocks) blocksize)))
           (output (make-blob outputsize)))
       (cbc-encrypt input 0 inputsize iv output 0 padfn encryptor blocksize))))

; As above, but the encryptor stores the IV in the output block (encrypted)...
(define (make-cbc*-encryptor encryptor padfn blocksize)
  (lambda (input iv)
     (let* ((inputsize (blob-size input))
            (output-whole-blocks (quotient inputsize blocksize))
            (output-overflow (remainder inputsize blocksize))
            (outputsize (if (zero? output-overflow) ; Round up to block size, plus an extra block for the IV
               (+ inputsize blocksize blocksize)
               (* (+ 2 output-whole-blocks) blocksize)))
            (output (make-blob outputsize))
            (encrypted-iv (encryptor iv)))
        (move-memory! encrypted-iv output blocksize)
        (cbc-encrypt input 0 inputsize iv output blocksize padfn encryptor blocksize))))

;internal cbc-decrypt function used by the make- variants. an unpad function is needed.
(define (cbc-decrypt input inoffset inputsize iv output outoffset unpadfn decryptor blocksize)
    (if (= (+ inoffset blocksize) inputsize)
       ; Last block
       (let* ((inblock (subblob input inoffset blocksize))
              (outblock
                 (unpadfn
                    (blob-xor iv
                      (decryptor inblock)))))
             (move-memory! outblock output (blob-size outblock) 0 outoffset)
             (subblob output 0 (+ outoffset (blob-size outblock)))) ; Terminate
       ; Not last block
       (let* ((inblock (subblob input inoffset blocksize))
              (outblock
                 (blob-xor iv
                   (decryptor inblock))))
             (move-memory! outblock output blocksize 0 outoffset)
             (cbc-decrypt input (+ inoffset blocksize) inputsize inblock output (+ outoffset blocksize) unpadfn decryptor blocksize)))) ; Recurse

;creates a cbc decryptor that uses and iv to start decrypting from the first byte of the input blob.
;an unpad function is needed in cbc mode decryption.
(define (make-cbc-decryptor decryptor unpadfn blocksize)
  (lambda (input iv)
    (let* ((inputsize (blob-size input))
           (output (make-blob inputsize)))
       (cbc-decrypt input 0 inputsize iv output 0 unpadfn decryptor blocksize))))

;same as above but treats the first block of the input as the encrypted IV.
(define (make-cbc*-decryptor decryptor unpadfn blocksize)
  (lambda (input)
    (let* ((inputsize (blob-size input))
           (output (make-blob inputsize))
           (iv (decryptor (subblob input 0 blocksize))))
       (cbc-decrypt input blocksize inputsize iv output 0 unpadfn decryptor blocksize)))) )