;; A Common Lisp arcfour implementation
;; Put into the Public Domain.

;; TODO:
;; - Write with-key-init or such macro
;; - Write stream-type functions for it


(declaim (optimize (speed 3) (safety 1))
	 (inline -key-init)
	 (inline key-random)
	 (inline key-random-add))


(defpackage :mm-rc4
  (:use :cl :sb-gray)
  (:nicknames :rc4)
  (:export
   #:rc4-type
   #:key-init
   #:key-random-add
   #:key-random
   #:key-encode
   #:key-destroy
   #:rc4-stream))


(in-package :rc4)


(defclass rc4-type ()
  ((x
    :type fixnum
    :accessor key-x
    :initform 0)
   (y
    :type fixnum
    :accessor key-y
    :initform 0)
   (state
    :type (simple-array fixnum (256))
    :accessor key-state
    :initform (make-array 256 :element-type 'fixnum :initial-element 0))))


(defgeneric -key-init (key))

(defmethod -key-init ((k rc4-type))
  (setf (key-x k) 0)
  (setf (key-y k) 0)
  (loop
     for i from 0 to 255
     do (setf (aref (the (simple-array fixnum (256)) (key-state k)) i) i)))


;; It is recommended to skip 3072 bytes.  Moreover, no two identical keys
;; must ever be used.
(defgeneric key-init (key data discard))

(defmethod key-init ((k rc4-type) (s string) (discard fixnum))
  (-key-init k)
  (key-random-add k s discard))

(defmethod key-init ((k rc4-type) (d simple-array (unsigned-byte 8) (*))
		     (discard fixnum))
  (-key-init k)
  (key-random-add k d discard))

(defmethod key-init ((k rc4-type) (l list) (discard fixnum))
  (-key-init k)
  (key-random-add k l discard))


(defgeneric key-random-add (key data discard))

(defmethod key-random-add ((k rc4-type) (s string) (discard fixnum))
  (let ((y 0) (si 0) (slen (length s)))
    (with-accessors ((sa key-state)) k
      (declare (type fixnum y si slen)
	       (type (simple-array fixnum (256)) sa)
	       (type (simple-array character (*)) s))
      (loop
	 for x from 0 to 255
	 do
	   (setf
	    y (logand (+ y (aref sa x) (char-code (char s si))) #Xff)
	    si (mod (1+ si) slen))
	   (psetf (aref sa x) (aref sa y)
		  (aref sa y) (aref sa x)))))
  (loop
     repeat discard
     do (key-random k)))

(defmethod key-random-add ((k rc4-type) (d simple-array (unsigned-byte 8) (*))
			   (discard fixnum))
  (let ((y 0) (di 0) (dlen (length d)))
    (with-accessors ((sa key-state)) k
      (declare (type fixnum y di dlen)
	       (type (simple-array fixnum (256)) sa)
	       (type (simple-array (unsigned-byte 8) (*)) d))
      (loop
	 for x from 0 to 255
	 do
	   (setf
	    y (logand (+ y (aref sa x) (the fixnum (aref d di))) #Xff)
	    di (mod (1+ di) dlen))
	   (psetf (aref sa x) (aref sa y)
		  (aref sa y) (aref sa x)))))
  (loop
     repeat discard
     do (key-random k)))

(defmethod key-random-add ((k rc4-type) (l list) (discard fixnum))
  (let ((slen (length l)) a)
    (declare (type fixnum slen)
	     (type (simple-array (unsigned-byte 8) (*))))
    (setf a (make-array slen :element-type '(unsigned-byte 8)))
    (loop
       for n in l
       for i from 0 to slen
       do (setf (aref a i) n))
    (key-random-add k a discard)))


(defgeneric key-random (key))

(defmethod key-random ((k rc4-type))
  (with-accessors ((sa key-state) (kx key-x) (ky key-y)) k
    (declare (type fixnum kx ky)
	     (type (simple-array fixnum (256)) sa))
    (setf
     kx (logand (1+ kx) #Xff)
     ky (logand (+ ky (aref sa kx)) #Xff))
    (psetf (aref sa kx) (aref sa ky)
	   (aref sa ky) (aref sa kx))
    (aref sa (logand (+ (aref sa kx) (aref sa ky)) #Xff))))


(defgeneric key-encode (k s))

;; Currently only supports destructive mode.
;; Note that using the unsigned-byte vector function is recommended
;; rather than the string one, however this is easier to use with REPL.
;; Also returns the string.
(defmethod key-encode ((k rc4-type) (s string))
  (let ((slen (max 0 (1- (length s)))))
    (declare (type fixnum slen)
	     (type (simple-array character (*)) s))
    (loop
       for i from 0 to slen
       do (setf (char s i)
		(code-char
		 (logxor (char-code (char s i))
			 (the fixnum (key-random k)))))))
  s)

(defmethod key-encode ((k rc4-type) (d simple-array (unsigned-byte 8) (*)))
  (let ((dlen (max 0 (1- (length d)))))
    (declare (type fixnum dlen)
	     (type (simple-array (unsigned-byte 8) (*)) d))
    (loop
       for i from 0 to dlen
       do (setf (aref d i)
	        (logxor (the fixnum (aref d i))
			(the fixnum (key-random k)))))))

(defmethod key-encode ((k rc4-type) (l list))
  (mapcar #'(lambda (x)
	      (logxor (the fixnum x)
		      (the fixnum (key-random k))))
	  l))


(defgeneric key-destroy (k))

(defmethod key-destroy ((k rc4-type))
  (setf (key-x k) 0)
  (setf (key-y k) 0)
  (loop
     for i from 0 to 255
     do (setf (aref (the (simple-array fixnum (256)) (key-state k)) i) 0)))


;; The following is incomplete and needs to be tested.

(defclass rc4-stream (fundamental-stream)
  ((state-output
    :initarg :key-in
    :initform nil
    :type rc4-type
    :reader stream-state-output)
   (state-input
    :initarg :key-out
    :initform nil
    :type rc4-type
    :reader stream-state-input)
   (stream
    :initarg :stream
    :reader stream-of)
   (buffer
    :type (simple-array (unsigned-byte 8) (*))
    :accessor stream-buffer
    :initform (make-array 4096 :element-type '(unsigned-byte 8)))))


(defmethod stream-element-type ((stream rc4-stream))
  '(unsigned-byte 8))

(defmethod close ((stream rc4-stream) &key (abort nil))
  (close (stream-of stream) :abort abort)
  (key-destroy (stream-state-input stream))
  (key-destroy (stream-state-output stream)))

(defclass rc4-output-stream (rc4-stream fundamental-binary-output-stream)
  ())

(defmethod stream-write-byte ((stream rc4-output-stream) b)
  (declare (type (unsigned-byte 8) b))
  (stream-write-byte (logand
		      (logxor (the fixnum b)
			      (the fixnum (key-random (stream-state-output
						       stream))))
		      #Xff)
		     (stream-of stream)))

(defmethod stream-write-sequence ((stream rc4-output-stream)
				  data &optional (start 0) (end nil))
  (let ((dlen (length data)))
    (with-accessors ((buf stream-buffer)) stream
      (declare (type (simple-array (unsigned-byte 8) (*)) data buf)
	       (type fixnum dlen))
    (if (< (length buf) dlen)
	(setf buf
	      (make-array dlen :element-type '(unsigned-byte 8))))
    ;;; gotta check start/end!
    (loop
       for i from 0 to dlen
       do
	 (setf (aref buf i)
	       (logand
		(logxor (the fixnum (aref data i))
			(the fixnum
			  (key-random (stream-state-output stream))))
		#Xff)))
    ;;; xxx take in consideration start/end
    (stream-write-sequence (stream-of stream) buf start end))))

(defclass rc4-input-stream (rc4-stream fundamental-binary-input-stream)
  ())

(defmethod stream-read-byte ((stream rc4-input-stream))
  (logand (logxor (the fixnum (stream-read-byte (stream-of stream)))
		  (the fixnum (key-random (stream-state-input stream))))
	  #Xff))

(defmethod stream-read-sequence ((stream rc4-input-stream)
				 data &optional (start 0) (end nil))
  (let ((dlen 0))
    (declare (type (simple-array (unsigned-byte 8) (*)) data)
	     (type fixnum dlen))
    ;;; xxx take care of start/end
    (stream-read-sequence (stream-of stream) data start end)
    (setf dlen (length data))
    (loop
       for i from 0 to dlen
       do
	 (setf (aref data i)
	       (logand
		(logxor (the fixnum (aref data i))
			(the fixnum
			  (key-random (stream-state-input stream))))
		#Xff)))))


;;; XXX using SBCL 1.0.16 X86
(defvar *k1* nil)
(defvar *k2* nil)
(defvar *s* nil)
(setf *k1* (make-instance 'rc4-type))
(setf *k2* (make-instance 'rc4-type))
(key-init *k1* "this is a test" 3072)
(key-init *k2* "this is a test" 3072)
(setf *s* (make-instance 'rc4-stream
			 :stream *standard-output*
			 :key-in *k1*
			 :key-out *k2*))
(stream-write-byte *s* 31)
(stream-write-sequence *s* #(12 32 53 64 13 43 4 24 53))

;; The stream #<RC4-STREAM {621E9821}> has no suitable method for
;; STREAM-WRITE-BYTE, and so has fallen through to this method.  If
;; you think that this is a bug, please report it to the applicable
;; authority (bugs in SBCL itself should go to the mailing lists
;; referenced from <http://www.sbcl.org/>).
;;   [Condition of type SIMPLE-ERROR]

;; There is no applicable method for the generic function
;;  #<STANDARD-GENERIC-FUNCTION STREAM-WRITE-SEQUENCE (3)>
;; when called with arguments
;;  (#<RC4-STREAM {62664121}> #(12 32 53 64 13 43 4 24 53)).
;;   [Condition of type SIMPLE-ERROR]