;;;; PMSF-Lib --- PMSF Common Lisp Utility Library
;;;; This is copyrighted software.  See documentation for terms.
;;;; 
;;;; aes-utilities.lisp --- Utilities concerned with AES
;;;; 
;;;; $Id$

(cl:in-package #:pmsf-lib)

;;;; %File Description:
;;;; 
;;;; This file contains utilities for AES encryption
;;;; 

(defvar *aes-processing-pathname* nil)

(define-condition aes-processing-error (simple-error)
  ((pathname :initarg :pathname :initform *aes-processing-pathname*
             :reader aes-processing-error-pathname))
  (:report 
   (lambda (c s)
     (with-standard-io-syntax
       (let ((*print-readably* nil))
         (format s 
                 "Error processing AES encryption~@[ for ~A~]: ~?"
                 (aes-processing-error-pathname c)
                 (simple-condition-format-control c)
                 (simple-condition-format-arguments c)))))))

(defun aes-processing-error (datum &rest arguments)
  (error 'aes-processing-error :format-control datum :format-arguments arguments))

(defparameter *crypt-aes-provider*
  "Microsoft Enhanced RSA and AES Cryptographic Provider")
(defconstant +prov-rsa-aes+ 24)
(defconstant +crypt-verify-context+ #xF0000000)
(defconstant +crypt-silent+ #x00000040)

(fli:define-foreign-function (crypt-acquire-context "CryptAcquireContext" :dbcs)
    ((:ignore (:reference-return win32:handle))
     (container (:reference-pass (:ef-wc-string :external-format :unicode :null-terminated-p t) :allow-null t))
     (provider (:reference-pass (:ef-wc-string :external-format :unicode :null-terminated-p t)))
     (prov-type win32:dword)
     (flags win32:dword))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-release-context "CryptReleaseContext")
    ((hprov win32:handle)
     (flags win32:dword))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-c-typedef crypt-aes-256-key-data
  (:c-array (:unsigned :byte) 32))

(fli:define-c-typedef crypt-aes-256-iv-data
  (:c-array (:unsigned :byte) 16))

(fli:define-c-struct crypt-aes-256-key
  (type (:unsigned :byte))
  (version (:unsigned :byte))
  (reserved win32:word)
  (alg-id (:unsigned :int))
  (key-size win32:dword)
  (key-data crypt-aes-256-key-data))

(fli:define-foreign-function (crypt-gen-random "CryptGenRandom")
    ((hprov win32:handle)
     (data-len win32:dword)
     (data (:pointer (:unsigned :byte))))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-import-key "CryptImportKey")
    ((hprov win32:handle)
     (data (:pointer crypt-aes-256-key))
     (data-len win32:dword)
     (pub-key win32:handle)
     (flags win32:dword)
     (:ignore (:reference-return win32:handle)))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-set-key-param "CryptSetKeyParam")
    ((hkey win32:handle)
     (param win32:dword)
     (data (:pointer (:unsigned :byte)))
     (flags win32:dword))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-encrypt "CryptEncrypt")
    ((hkey win32:handle)
     (hhash win32:handle)
     (finalp win32:bool)
     (flags win32:dword)
     (data (:pointer (:unsigned :byte)))
     (data-len (:reference win32:dword))
     (buf-len win32:dword))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-decrypt "CryptDecrypt")
    ((hkey win32:handle)
     (hhash win32:handle)
     (finalp win32:bool)
     (flags win32:dword)
     (data (:pointer (:unsigned :byte)))
     (data-len (:reference win32:dword)))
  :result-type win32:bool
  :calling-convention :stdcall)

(fli:define-foreign-function (crypt-destroy-key "CryptDestroyKey")
    ((hkey win32:handle))
  :result-type win32:bool
  :calling-convention :stdcall)

(defmacro with-aes-winapi-call (vars form cleanup &body body)
  (with-unique-names (ok-p)
    `(multiple-value-bind (,ok-p ,@vars) ,form
       (unless ,ok-p
         (aes-processing-error
          "Failed CryptAPI call (~A), Error: ~A" ',(car form)
          (win32:get-last-error-string)))
       (unwind-protect (progn ,@body)
         ,cleanup))))

(defun aes-encode-buffer (buffer size &key key-in iv-in)
  (with-aes-winapi-call (hprov) 
      (crypt-acquire-context nil *crypt-aes-provider*
                             +prov-rsa-aes+
                             (logior +crypt-verify-context+
                                     +crypt-silent+))
      (crypt-release-context hprov 0)
    (fli:with-dynamic-foreign-objects ((key crypt-aes-256-key)
                                       (iv crypt-aes-256-iv-data))
      (setf (fli:foreign-slot-value key 'type) #x8
            (fli:foreign-slot-value key 'version) 2
            (fli:foreign-slot-value key 'reserved) 0
            (fli:foreign-slot-value key 'alg-id) #x00006610
            (fli:foreign-slot-value key 'key-size) 32)
      (let ((key-data (fli:foreign-slot-pointer key 'key-data)))
        (if key-in
            (loop for i upfrom 0 for val across key-in
                  do (setf (fli:foreign-aref key-data i) val))
          (with-aes-winapi-call ()
              (crypt-gen-random hprov 32 (fli:foreign-array-pointer key-data 0))
              t)))
      (if iv-in
          (loop for i upfrom 0 for val across iv-in
                do (setf (fli:foreign-aref iv i) val))
        (with-aes-winapi-call ()
            (crypt-gen-random hprov 16 (fli:foreign-array-pointer iv 0))
            t))
      (with-aes-winapi-call (hkey)
          (crypt-import-key hprov key (+ 32 4 4 4) 0 0)
          (crypt-destroy-key hkey)
        (with-aes-winapi-call ()
            (crypt-set-key-param hkey 1 (fli:foreign-array-pointer iv 0) 0)
            t
          (fli:with-dynamic-lisp-array-pointer (buffer-ptr buffer :type '(:unsigned :byte))
            (with-aes-winapi-call (enc-len)
                (crypt-encrypt hkey 0 t 0 buffer-ptr size (length buffer))
                t
              (values
               buffer
               enc-len
               (let ((key-out (make-array 32 :element-type '(unsigned-byte 8))))
                 (fli:replace-foreign-array key-out (fli:foreign-slot-pointer key 'key-data))
                 key-out)
               (let ((iv-out (make-array 16 :element-type '(unsigned-byte 8))))
                 (fli:replace-foreign-array iv-out iv)
                 iv-out)))))))))

(defun aes-decode-buffer (buffer size key-in iv-in)
  (with-aes-winapi-call (hprov) 
      (crypt-acquire-context nil *crypt-aes-provider*
                             +prov-rsa-aes+
                             (logior +crypt-verify-context+
                                     +crypt-silent+))
      (crypt-release-context hprov 0)
    (fli:with-dynamic-foreign-objects ((key crypt-aes-256-key)
                                       (iv crypt-aes-256-iv-data))
      (setf (fli:foreign-slot-value key 'type) #x8
            (fli:foreign-slot-value key 'version) 2
            (fli:foreign-slot-value key 'reserved) 0
            (fli:foreign-slot-value key 'alg-id) #x00006610
            (fli:foreign-slot-value key 'key-size) 32)
      (let ((key-data (fli:foreign-slot-pointer key 'key-data)))
        (loop for i upfrom 0 for val across key-in
              do (setf (fli:foreign-aref key-data i) val)))
      (loop for i upfrom 0 for val across iv-in
            do (setf (fli:foreign-aref iv i) val))
      (with-aes-winapi-call (hkey)
          (crypt-import-key hprov key (+ 32 4 4 4) 0 0)
          (crypt-destroy-key hkey)
        (with-aes-winapi-call ()
            (crypt-set-key-param hkey 1 (fli:foreign-array-pointer iv 0) 0)
            t
          (fli:with-dynamic-lisp-array-pointer (buffer-ptr buffer :type '(:unsigned :byte))
            (with-aes-winapi-call (dec-len)
                (crypt-decrypt hkey 0 t 0 buffer-ptr size)
                t
              (values
               buffer
               dec-len))))))))

(defun aes-encode-file (from to &key key-in iv-in)
  (let ((*aes-processing-pathname* from))
    (with-open-file (in from :element-type '(unsigned-byte 8))
      (with-open-file (out to :element-type '(unsigned-byte 8)
                           :direction :output)
        (let* ((length (file-length in))
               (buffer (make-array (* (ceiling (1+ length) 16) 16)
                                   :element-type '(unsigned-byte 8)
                                   :allocation :static)))
          (unless (= length (read-sequence buffer in :end length))
            (aes-processing-error 
             "Failed to read ~D bytes from ~A!" length from))
          (multiple-value-bind (out-buffer out-length key iv)
              (aes-encode-buffer buffer length 
                                 :key-in key-in :iv-in iv-in)
            (write-sequence out-buffer out :end out-length)
            (values length key iv)))))))

(defun aes-decode-file (from to key-in iv-in)
  (let ((*aes-processing-pathname* from))
    (with-open-file (in from :element-type '(unsigned-byte 8))
      (with-open-file (out to :element-type '(unsigned-byte 8)
                           :direction :output)
        (let* ((length (file-length in))
               (buffer (make-array length
                                   :element-type '(unsigned-byte 8)
                                   :allocation :static)))
          (unless (= length (read-sequence buffer in :end length))
            (aes-processing-error
             "Failed to read ~D bytes from ~A!" length from))
          (multiple-value-bind (out-buffer out-length)
              (aes-decode-buffer buffer length key-in iv-in) 
            (write-sequence out-buffer out  :end out-length)
            out-length))))))