; elsie four

(define (drop n xs)
  (let loop ((n n) (xs xs))
    (if (or (zero? n) (null? xs)) xs
      (loop (- n 1) (cdr xs)))))

(define (make-matrix rows columns . value)
  (do ((m (make-vector rows)) (i 0 (+ i 1)))
      ((= i rows) m)
    (if (null? value)
        (vector-set! m i (make-vector columns))
        (vector-set! m i (make-vector columns (car value))))))

(define (matrix-ref m i j) (vector-ref (vector-ref m i) j))

(define (matrix-set! m i j x) (vector-set! (vector-ref m i) j x))

(define (string->ints str)
  (define (c->i c)
    (cond ((char=? c #\#) 0) ((char=? c #\_) 1)
          ((char<=? #\2 c #\9) (- (char->integer c) (char->integer #\0)))
          ((char<=? #\a c #\z) (- (char->integer c) (char->integer #\a) -10))
          ((char<=? #\A c #\Z) (- (char->integer c) (char->integer #\A) -10))
          (else (error 'string->ints "unrecognized character"))))
  (map c->i (string->list str)))

(define (ints->string ints)
  (define (i->c i)
    (cond ((or (< i 0) (< 35 i)) (error 'ints->string "unrecognized character"))
          ((= i 0) #\#) ((= i 1) #\_)
          ((< i 10) (integer->char (+ (char->integer #\0) i)))
          (else (integer->char (+ (char->integer #\A) i -10)))))
  (list->string (map i->c ints)))

(define (initialize-state key)
  (let ((xs (string->ints key)) (state (make-matrix 6 6)))
    (do ((k 0 (+ k 1)) (xs xs (cdr xs))) ((null? xs) state)
      (matrix-set! state (quotient k 6) (modulo k 6) (car xs)))))

(define (find s x)
  (let loop ((k 0))
    (let ((r (quotient k 6)) (c (modulo k 6)))
      (if (= (matrix-ref s r c) x) (values r c)
        (loop (+ k 1))))))

(define (advance-state state i j r c x y ct)
  (let ((t (matrix-ref state r 5)))
    (matrix-set! state r 5 (matrix-ref state r 4))
    (matrix-set! state r 4 (matrix-ref state r 3))
    (matrix-set! state r 3 (matrix-ref state r 2))
    (matrix-set! state r 2 (matrix-ref state r 1))
    (matrix-set! state r 1 (matrix-ref state r 0))
    (matrix-set! state r 0 t))
  (set! c (modulo (+ c 1) 6))
  (when (= x r) (set! y (modulo (+ y 1) 6)))
  (when (= i r) (set! j (modulo (+ j 1) 6)))
  (let ((t (matrix-ref state 5 y)))
    (matrix-set! state 5 y (matrix-ref state 4 y))
    (matrix-set! state 4 y (matrix-ref state 3 y))
    (matrix-set! state 3 y (matrix-ref state 2 y))
    (matrix-set! state 2 y (matrix-ref state 1 y))
    (matrix-set! state 1 y (matrix-ref state 0 y))
    (matrix-set! state 0 y t))
  (set! x (modulo (+ x 1) 6))
  (when (= c y) (set! r (modulo (+ r 1) 6)))
  (when (= j y) (set! i (modulo (+ i 1) 6)))
  (set! i (modulo (+ i (quotient ct 6)) 6))
  (set! j (modulo (+ j (modulo ct 6)) 6))
  (values state i j r c x y))

(define (encrypt key nonce plaintext signature)
  (let ((state (initialize-state key)) (i 0) (j 0))
    (let loop ((ps (string->ints (string-append nonce plaintext signature)))
               (state state) (i i) (j j) (cs (list)))
      (if (null? ps)
          (string-upcase (string-append nonce
            (ints->string (drop (string-length nonce) (reverse cs)))))
          (call-with-values
            (lambda () (find state (car ps)))
            (lambda (r c)
              (let* ((x (modulo (+ r (quotient (matrix-ref state i j) 6)) 6))
                     (y (modulo (+ c (modulo (matrix-ref state i j) 6)) 6))
                     (ct (matrix-ref state x y)))
                (call-with-values
                  (lambda () (advance-state state i j r c x y ct))
                  (lambda (state i j r c x y)
                    (loop (cdr ps) state i j (cons ct cs)))))))))))

(define (decrypt key nonce ciphertext)
  (let ((state (initialize-state key)) (i 0) (j 0))
    ; encrypt nonce (to set initial state for decryption)
    (let loop ((ps (string->ints nonce)) (state state) (i i) (j j))
      (when (pair? ps)
        (call-with-values
          (lambda () (find state (car ps)))
          (lambda (r c)
            (let* ((x (modulo (+ r (quotient (matrix-ref state i j) 6)) 6))
                   (y (modulo (+ c (modulo (matrix-ref state i j) 6)) 6))
                   (ct (matrix-ref state x y)))
              (call-with-values
                (lambda () (advance-state state i j r c x y ct))
                (lambda (state i j r c x y)
                  (loop (cdr ps) state i j))))))))
    ; decrypt message text and signature
    (let loop ((cs (string->ints ciphertext)) (state state) (i i) (j j) (ps (list)))
      (if (null? cs)
          (string-upcase (ints->string (reverse ps)))
          (call-with-values
            (lambda () (find state (car cs)))
            (lambda (x y)
              (let* ((r (modulo (- x (quotient (matrix-ref state i j) 6)) 6))
                     (c (modulo (- y (modulo (matrix-ref state i j) 6)) 6))
                     (pt (matrix-ref state r c)))
                (call-with-values
                  (lambda () (advance-state state i j r c x y (car cs)))
                  (lambda (state i j r c x y)
                    (loop (cdr cs) state i j (cons pt ps)))))))))))

(define key "XV7YDQ#OPAJ_39RZUT8B45WCSGEHMIKNF26L")
(define nonce "SOLWBF")
(define plaintext "IM_ABOUT_TO_PUT_THE_HAMMER_DOWN")
(define signature "#RUBBERDUCK")

(display (encrypt key nonce plaintext signature)) (newline)
(display (decrypt key "SOLWBF" "I2ZQPILR2YQGPTLTRZX2_9FZLMBO3Y8_9PYSSX8NF2")) (newline)