;;;; cipher.lisp -- generic functions for symmetric encryption

(in-package :crypto)

(defgeneric encrypt (context plaintext ciphertext
                             &key plaintext-start plaintext-end
                                  ciphertext-start)
  (:documentation "Encrypt the data in PLAINTEXT between PLAINTEXT-START
and PLAINTEXT-END according to CONTEXT.  Places the encrypted data in
CIPHERTEXT, beginning at CIPHERTEXT-START.  Less data than
(- PLAINTEXT-END PLAINTEXT-START) may be encrypted, depending on the
alignment constraints of CONTEXT and the amount of space available in
CIPHERTEXT."))

(defgeneric decrypt (context ciphertext plaintext
                             &key ciphertext-start ciphertext-end
                                  plaintext-start)
  (:documentation "Decrypt the data in CIPHERTEXT between CIPHERTEXT-START
and CIPHERTEXT-END according to CONTEXT.  Places the decrypted data in
PLAINTEXT, beginning at PLAINTEXT-START.  Less data than
(- CIPHERTEXT-END CIPHERTEXT-START) may be encrypted, depending on the
alignment constraints of CONTEXT and the amount of psace available in
PLAINTEXT."))

(defun encrypt-in-place (context text &key (start 0) end)
  (encrypt context text text
           :plaintext-start start :plaintext-end (or end (length text))
           :ciphertext-start start))

(defun decrypt-in-place (context text &key (start 0) end)
  (decrypt context text text
           :ciphertext-start start :ciphertext-end (or end (length text))
           :plaintext-start start))


;;; utilities for wordwise fetches and stores

;;; we attempt to make this as efficient as possible.  the first check we
;;; do is to see whether or not the range
;;; [INITIAL-OFFSET, INITIAL-OFFSET + BLOCK-SIZE) is within the bounds of
;;; the array.  if not, then we do the fetches as normal.  if so, then we
;;; can either (DECLARE (SAFETY 0)) to avoid the bounds-checking on the
;;; fetches, or we can do full-word fetches if INITIAL-OFFSET is
;;; word-addressable and the implementation supports it.
#+nil
(defmacro with-words (((&rest word-vars) array initial-offset
                       &key (size 4) (big-endian t))
                      &body body)
  (let ((ref-sym (intern (format nil "UB~DREF/~:[LE~;BE~]"
                                 (* size 8) big-endian)))
        (n-bytes (* (length word-vars) size)))
    (flet ((generate-fetches (n-fetches)
             (loop for offset from 0 by size below (* n-fetches size)
                   collect `(,ref-sym ,array (+ ,initial-offset ,offset)))))
      `(multiple-value-bind ,word-vars (let ((length (length ,array)))
                                         (cond
                                           ((<= ,initial-offset (- length ,n-bytes))
                                            ,(if (and (member :sbcl *features*)
                                                      (= size 4)
                                                      (or (and big-endian (member :big-endian *features*))
                                                          (and (not big-endian) (member :little-endian *features*))))
                                                 `(if (logtest ,initial-offset (1- ,size))
                                                      ;; do FETCH-UB* way
                                                      (locally (declare (optimize (safety 0)))
                                                        (values ,@(generate-fetches (length word-vars))))
                                                      (let ((word-offset (truncate ,initial-offset 4)))
                                                        (values
                                                         ,@(loop for offset from 0 below (length word-vars)
                                                              collect `(sb-kernel:%vector-raw-bits ,array (+ word-offset ,offset))))))
                                                  `(locally (declare (optimize (safety 0)))
                                                     (values ,@(generate-fetches (length word-vars))))))
                                           (t
                                            (values ,@(generate-fetches (length word-vars))))))
         (declare (type (unsigned-byte ,(* size 8)) ,@word-vars))
         (macrolet ((store-words (buffer buffer-offset &rest word-vars)
                      (loop for word-var in word-vars
                         for offset from 0 by ,size
                         collect `(setf (,',ref-sym ,buffer (+ ,buffer-offset ,offset)) ,word-var)
                         into stores
                         finally (return `(progn ,@stores)))))
           ,@body)))))

(defmacro with-words (((&rest word-vars) array initial-offset
                       &key (size 4) (big-endian t))
                      &body body)
  (let ((ref-sym (intern (format nil "UB~DREF/~:[LE~;BE~]"
                                   (* size 8) big-endian))))
    (loop for word-var in word-vars
          for offset from 0 by size
          collect `(,word-var (,ref-sym ,array (+ ,initial-offset ,offset)))
          into let-bindings
          finally (return `(macrolet ((store-words (buffer buffer-offset &rest word-vars)
                                       (loop for word-var in word-vars
                                             for offset from 0 by ,size
                                             collect `(setf (,',ref-sym ,buffer (+ ,buffer-offset ,offset)) ,word-var)
                                             into stores
                                             finally (return `(progn ,@stores)))))
                             (let ,let-bindings
                               (declare (type (unsigned-byte ,(* size 8)) ,@word-vars))
                               ,@body))))))



;;; defining ciphers

;;; the idea behind this is that one only has to implement encryption
;;; and decryption of a block for a particular cipher (and perhaps
;;; some key generation) and then "define" the cipher with some
;;; parameters.  necessary interface functions will be auto-generated
;;; with this macro.

;;; possible things to go in INITARGS
;;;
;;; * (:encrypt-function #'cipher-encrypt-block)
;;; * (:decrypt-function #'cipher-decrypt-block)
;;; * (:block-length number)
;;; * (:key-length (:fixed &rest lengths))
;;; * (:key-length (:variable low high increment))
;;; * (:constructor #'create-cipher-context)
(defmacro defcipher (name &rest initargs)
  (%defcipher name initargs))

;;; KLUDGE: we add the blocksize to these two forms so that we can declare
;;; the type of the *-START parameters correctly.  That is, good Lisp
;;; implementations will see that references into the plaintext and
;;; ciphertext can never overflow into bignum land; shorter code should
;;; then be generated.  This is a kludge, however, because we're putting
;;; the blocksize in three different places: once in the encryptor, once
;;; in the decryptor, and once in the DEFCIPHER form.  It would be nice
;;; if there was one single place to put everything.
(defmacro define-block-encryptor (algorithm blocksize &body body)
  `(defun ,(intern (format nil "~A-ENCRYPT-BLOCK" algorithm))
    (context plaintext plaintext-start ciphertext ciphertext-start)
    (declare (optimize (speed 3) (debug 0) (space 0)))
    (declare (type (simple-array (unsigned-byte 8) (*)) plaintext ciphertext)
     (type (integer 0 ,(- array-dimension-limit blocksize))
      plaintext-start ciphertext-start))
    ,@body))

(defmacro define-block-decryptor (algorithm blocksize &body body)
  `(defun ,(intern (format nil "~A-DECRYPT-BLOCK" algorithm))
    (context ciphertext ciphertext-start plaintext plaintext-start)
    (declare (optimize (speed 3) (debug 0) (space 0)))
    (declare (type (simple-array (unsigned-byte 8) (*)) ciphertext plaintext)
     (type (integer 0 ,(- array-dimension-limit blocksize))
      ciphertext-start plaintext-start))
    ,@body))

(defmacro define-stream-cryptor (algorithm &body body)
  `(defun ,(intern (format nil "~A-CRYPT" algorithm))
       (context plaintext plaintext-start ciphertext ciphertext-start length)
     (declare (optimize (speed 3) (debug 0) (space 0)))
     (declare (type (simple-array (unsigned-byte 8) (*)) plaintext ciphertext))
     (declare (type index plaintext-start ciphertext-start length))
     ,@body))

(defgeneric verify-key (cipher key)
  (:documentation "Return T is KEY is a valid encryption key for CIPHER."))

(defgeneric block-length (cipher)
  (:documentation "Return the number of bytes in an encryption or decryption
block for CIPHER."))

(defgeneric key-lengths (cipher)
  (:documentation "Return a list of possible lengths of a key for CIPHER."))

;;; introspection
(defvar *supported-ciphers* nil)

(defun list-all-ciphers ()
  (copy-seq *supported-ciphers*))

(defun cipher-supported-p (name)
  "Return T if the cipher NAME is supported as an argument to MAKE-CIPHER."
  (member name *supported-ciphers*))

(defgeneric make-cipher (name mode key &optional iv)
  (:documentation "Return a cipher object using algorithm NAME with KEY in the
specified MODE.  If MODE requires an initialization vector, it must be provided
as IV; otherwise, the IV argument is ignored."))

(defmethod make-cipher (name mode key &optional iv)
  (error 'unsupported-cipher :name name))

(defun acceptable-key-lengths (key-length-spec)
  (ecase (car key-length-spec)
    (:fixed (cdr key-length-spec))
    (:variable (destructuring-bind (low high increment) (cdr key-length-spec)
                 (loop for i from low to high by increment
                       collect i)))))

(defun generate-key-verifier-methods (name key-length-spec)
  (let ((acceptable-key-lengths (acceptable-key-lengths key-length-spec)))
    `(defmethod verify-key ((cipher (eql ,name)) (key vector))
      (unless (equal (array-element-type key) '(unsigned-byte 8))
        (error 'type-error :expected-type '(vector (unsigned-byte 8))
               :datum key))
      (case (length key)
        (,acceptable-key-lengths (copy-seq key))
        (t (error 'invalid-key-length
                  :cipher ',name
                  :accepted-lengths ',acceptable-key-lengths))))))

(defun wrap-cipher-context-with-mode (context cipher-name 
                                      mode &optional initialization-vector)
  (let ((block-length (block-length context)))
    (case mode
      (:ecb
       (make-instance 'ecb-mode :cipher context))
      ((:cbc :ofb :cfb :ctr)
       (unless initialization-vector
         (error 'initialization-vector-not-supplied
                :mode mode))
       (unless (typep initialization-vector '(vector (unsigned-byte 8)))
         (error 'type-error
                :datum initialization-vector
                :expected-type '(vector (unsigned-byte 8))))
       (unless (= (length initialization-vector) block-length)
         (error 'invalid-initialization-vector
                :cipher cipher-name
                :block-length block-length))
       (make-instance (case mode
                        (:cbc 'cbc-mode)
                        (:ofb 'ofb-mode)
                        (:cfb 'cfb-mode)
                        (:ctr 'ctr-mode))
                      :initialization-vector (copy-seq initialization-vector)
                      :cipher context))
      (:stream context)
      (t
       (error 'unsupported-mode :mode mode)))))

(defun generate-make-cipher-defmethod (cipher-name constructor block-length
                                       &optional (block-cipher-p t))
  `(defmethod make-cipher ((cipher (eql ,cipher-name)) mode key
                           &optional initialization-vector)
    (declare (ignorable initialization-vector))
    ,(if block-cipher-p
         `(unless (member mode '(:cbc :ofb :cfb :ctr :ecb))
            (error 'unsupported-mode :mode mode))
         `(unless (eq mode :stream)
            (error 'unsupported-mode :mode mode)))
    (let* ((byte-key (verify-key ,cipher-name key))
           (cipher-context (,constructor byte-key)))
      (wrap-cipher-context-with-mode cipher-context ',cipher-name
                                     mode initialization-vector))))

(defun generate-common-cipher-methods (name block-length key-length-spec)
  (let ((context-name (intern (format nil "~A-CONTEXT" (symbol-name name))))
        (cipher-name (intern (symbol-name name) (find-package :keyword))))
    `(progn
       ;; make sure we pass in valid keys
       ,(generate-key-verifier-methods cipher-name key-length-spec)
       (pushnew ,cipher-name *supported-ciphers*)
       (defmethod block-length ((cipher (eql ,cipher-name)))
         ,block-length)
       (defmethod block-length ((cipher ,context-name))
         ,block-length)
       (defmethod key-lengths ((cipher (eql ,cipher-name)))
         (list ,@(acceptable-key-lengths key-length-spec)))
       (defmethod key-lengths ((cipher ,context-name))
         (list ,@(acceptable-key-lengths key-length-spec))))))

(defun generate-block-cipher-forms (name block-length key-length-spec
                                   encrypt-function decrypt-function)
  (declare (ignorable key-length-spec))
  (let ((constructor (intern (format nil "CREATE-~A-CONTEXT" (symbol-name name))))
        (context-name (intern (format nil "~A-CONTEXT" (symbol-name name))))
        (cipher-name (intern (symbol-name name) (find-package :keyword))))
    ;; ensure that ciphers can be created properly
    (loop for mode in '(:ecb :cbc :ofb :cfb :ctr)
          append (generate-cipher-mode-functions mode context-name
                                                 block-length
                                                 encrypt-function
                                                 decrypt-function) into forms
          ;; finally, all the work is done!  return all the cool stuff we define
          finally (return `(progn
                             ,(generate-make-cipher-defmethod cipher-name constructor
                                                              block-length)
                             ,@forms)))))

(defun generate-stream-cipher-forms (name key-length-spec crypt-function)
  (declare (ignorable key-length-spec))
  (let ((constructor (intern (format nil "CREATE-~A-CONTEXT" (symbol-name name))))
        (context-name (intern (format nil "~A-CONTEXT" (symbol-name name))))
        (cipher-name (intern (symbol-name name) (find-package :keyword))))
    `(progn
       ,@(generate-cipher-mode-functions :stream context-name 1
                                         crypt-function crypt-function)
       ,(generate-make-cipher-defmethod cipher-name constructor 1 nil))))

(defun %defcipher (name initargs)
  (let ((encrypt-function nil)
        (decrypt-function nil)
        (crypt-function nil)
        (mode :block)
        (block-length nil)
        (key-length-spec nil)
        (constructor nil))
    (declare (ignorable constructor))
    (loop for (arg value) in initargs
          do (case arg
               (:encrypt-function
                (if (not encrypt-function)
                    (setf encrypt-function value)
                    (error "Specified :ENCRYPT-FUNCTION multiple times.")))
               (:decrypt-function
                (if (not decrypt-function)
                    (setf decrypt-function value)
                    (error "Specified :DECRYPT-FUNCTION multiple times.")))
               (:crypt-function
                (if (not crypt-function)
                    (setf crypt-function value)
                    (error "Specified :CRYPT-FUNCTION multiple times.")))
               (:mode
                (setf mode value))
               (:block-length
                (cond
                  (block-length
                   (error "Specified :BLOCK-LENGTH multiple times."))
                  ((or (not (integerp value))
                       (not (plusp value)))
                   (error ":BLOCK-LENGTH must be a positive, integral number."))
                  (t
                   (setf block-length value))))
               (:key-length
                (cond
                  (key-length-spec
                   (error "Specified :KEY-LENGTH multiple times."))
                  ((not (consp value))
                   (error ":KEY-LENGTH value must be a list."))
                  ((and (not (eq :fixed (car value)))
                        (not (eq :variable (car value))))
                   (error "First element of :KEY-LENGTH spec must be either :FIXED or :VARIABLE."))
                  ((eq :fixed (car value))
                   (if (and (cdr value)
                            (every #'integerp (cdr value))
                            (every #'plusp (cdr value)))
                       (setf key-length-spec value)
                       ;;; FIXME: better error message
                       (error "bad :FIXED specification for :KEY-LENGTH.")))
                  ((eq :variable (car value))
                   (if (and (null (nthcdr 4 value))
                            (every #'integerp (cdr value))
                            (every #'plusp (cdr value))
                            (< (cadr value) (caddr value)))
                       (setf key-length-spec value)
                       (error "bad :VARIABLE specification for :KEY-LENGTH."))))))
          finally (cond
                    ((and (eq mode :block) block-length key-length-spec encrypt-function decrypt-function)
                     (return
                       `(progn
                          ,(generate-common-cipher-methods name block-length key-length-spec)
                          ,(generate-block-cipher-forms name block-length key-length-spec
                                                        encrypt-function decrypt-function))))
                    ((and (eq mode :stream) crypt-function key-length-spec)
                     (return
                       `(progn
                          ,(generate-common-cipher-methods name 1 key-length-spec)
                          ,(generate-stream-cipher-forms name key-length-spec crypt-function))))
                    (t
                     (error "Didn't specify all required fields for DEFCIPHER"))))))
