;;;; -*- mode: lisp; indent-tabs-mode: nil -*-
;;;; keystream.lisp

(in-package :crypto)


(defun keystream-position (cipher &optional position)
  "Return or change the current POSITION within the key stream of a CIPHER."
  (let ((mode (mode cipher)))
    (cond
      ((typep mode 'ctr-mode)
       (let ((block-length (block-length cipher))
             (iv-position (iv-position mode))
             (keystream-blocks (keystream-blocks mode)))
         (if (null position)
             (if (zerop iv-position)
                 (* block-length keystream-blocks)
                 (+ (* block-length (1- keystream-blocks)) iv-position))
             (let ((iv (iv mode))
                   (buffer (make-array block-length :element-type '(unsigned-byte 8))))
               (multiple-value-bind (q r)
                   (truncate position block-length)
                 (if (< q keystream-blocks)
                     (decrement-counter-block iv (- keystream-blocks q))
                     (increment-counter-block iv (- q keystream-blocks)))
                 (setf (keystream-blocks mode) q)
                 (setf (iv-position mode) 0)
                 (encrypt-in-place cipher buffer :end r)
                 t)))))

      ((typep cipher 'chacha)
       (let ((state (chacha-state cipher)))
         (if (null position)
             (let ((counter (if (= (chacha-counter-size cipher) 1)
                                (aref state 12)
                                (+ (aref state 12) (ash (aref state 13) 32)))))
               (- (* 64 counter) (chacha-keystream-buffer-remaining cipher)))
             (let ((buffer (make-array 64 :element-type '(unsigned-byte 8))))
               (declare (dynamic-extent buffer))
               (multiple-value-bind (q r)
                   (truncate position 64)
                 (setf (aref state 12) (logand q #xffffffff))
                 (unless (= (chacha-counter-size cipher) 1)
                   (setf (aref state 13) (logand (ash q -32) #xffffffff)))
                 (setf (chacha-keystream-buffer-remaining cipher) 0)
                 (encrypt-in-place cipher buffer :end r)
                 t)))))

      ((typep cipher 'salsa20)
       (let ((state (salsa20-state cipher)))
         (if (null position)
             (let ((counter (+ (aref state 8) (ash (aref state 9) 32))))
               (- (* 64 counter) (salsa20-keystream-buffer-remaining cipher)))
             (let ((buffer (make-array 64 :element-type '(unsigned-byte 8))))
               (declare (dynamic-extent buffer))
               (multiple-value-bind (q r)
                   (truncate position 64)
                 (setf (aref state 8) (logand q #xffffffff))
                 (setf (aref state 9) (logand (ash q -32) #xffffffff))
                 (setf (salsa20-keystream-buffer-remaining cipher) 0)
                 (encrypt-in-place cipher buffer :end r)
                 t)))))

      (t
       nil))))
