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

(import scheme)
(import chicken foreign)
(use extras lolevel)

(define *the-null-blob* (make-blob 0))

(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))))

(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)))

(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)))

(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))))))

(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))))

(define (blob-xor a b)
   (if (not (blob? a)) (error "blob-xor: Arguments must be blobs of the same size"))
   (check-blob b (blob-size a) 'blob-xor)
   
   (let ((out (make-blob (blob-size a))))
      ((foreign-lambda* void ((int len) (blob a) (blob b) (blob out)) "
         while (len--) {
            (*out++) = (*a++) ^ (*b++);
         }
      ") (blob-size a) a b out)
      out))

(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))

(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)))

(define (make-cbc-encryptor encryptor blocksize)
   (letrec ((encrypt (lambda (input inoffset inputsize iv output outoffset)
         (cond
            ((= inoffset inputsize) ; Zero bytes
               (let* ((inblock (blob-pad *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)
                  (encrypt input (+ inoffset blocksize) inputsize outblock output (+ outoffset blocksize)))) ; Recurse
            (else ; Partial block
               (let* ((inblock (blob-pad
                        (subblob input inoffset (- inputsize inoffset))
                        blocksize))
                      (outblock (encryptor (blob-xor iv inblock))))

                  (move-memory! outblock output blocksize 0 outoffset)
                  output)))))) ; Terminate

            (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)))
                  (encrypt input 0 inputsize iv output 0)))))

(define (make-cbc-decryptor decryptor blocksize)
   (letrec ((decrypt (lambda (input inoffset inputsize iv output outoffset)
               (if (= (+ inoffset blocksize) inputsize)
                  ; Last block
                  (let* ((inblock (subblob input inoffset blocksize))
                         (outblock
                            (blob-unpad
                               (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)
                        (decrypt input (+ inoffset blocksize) inputsize inblock output (+ outoffset blocksize))))))) ; Recurse

            (lambda (input iv)
               (let* ((inputsize (blob-size input))
                      (output (make-blob inputsize)))
                  (decrypt input 0 inputsize iv output 0)))))

; As above, but the encryptor stores the IV in the output block (encrypted)...

(define (make-cbc*-encryptor encryptor blocksize)
   (letrec ((encrypt (lambda (input inoffset inputsize iv output outoffset)
         (cond
            ((= inoffset inputsize) ; Zero bytes
               (let* ((inblock (blob-pad *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)
                  (encrypt input (+ inoffset blocksize) inputsize outblock output (+ outoffset blocksize)))) ; Recurse
            (else ; Partial block
               (let* ((inblock (blob-pad
                        (subblob input inoffset (- inputsize inoffset))
                        blocksize))
                      (outblock (encryptor (blob-xor iv inblock))))

                  (move-memory! outblock output blocksize 0 outoffset)
                  output)))))) ; Terminate

            (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)
                  (encrypt input 0 inputsize iv output blocksize)))))

;... and the decryptor retreives it.

(define (make-cbc*-decryptor decryptor blocksize)
   (letrec ((decrypt (lambda (input inoffset inputsize iv output outoffset)
               (if (= (+ inoffset blocksize) inputsize)
                  ; Last block
                  (let* ((inblock (subblob input inoffset blocksize))
                         (outblock
                            (blob-unpad
                               (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)
                        (decrypt input (+ inoffset blocksize) inputsize inblock output (+ outoffset blocksize))))))) ; Recurse

            (lambda (input)
               (let* ((inputsize (blob-size input))
                      (output (make-blob inputsize))
                      (iv (decryptor (subblob input 0 blocksize))))
                  (decrypt input blocksize inputsize iv output 0))))) )