enigma2.0

#lang racket
;; a rotor will be a list of 1) the settings 2) the numbers of times
;; it has rotated % 26 3) it's notch 4) the corresponding encryption
(define alph (string->list "ABCDEFGHIJKLMNOPQRSTUVWXYZ"))

(define rotor1 (string->list "EKMFLGDQVZNTOWYHXUSPAIBRCJ"))
(define rotor2 (string->list "AJDKSIRUXBLHWTMCQGZNPYFVOE"))
(define rotor3 (string->list "BDFHJLCPRTXVZNYEIWGAKMUSQO"))
(define rotor4 (string->list "ESOVPZJAYQUIRHXLNFTGKDCMWB"))
(define rotor5 (string->list "VZBRGITYUPSDNHLXAWMJQOFECK"))

(define refB (string->list "YRUHQSLDPXNGOKMIEBFZCWVJAT"))
(define refC (string->list "FVPJIAOYEDRZXWGCTKUQSBNMHL"))

(define rI (list alph 0 (string->list "Q") rotor1))
(define rII (list alph 0 (string->list "E") rotor2))
(define rIII (list alph 0 (string->list "V") rotor3))
(define rIV (list alph 0 (string->list "J") rotor4))
(define rV (list alph 0 (string->list "Z") rotor5))
(define rB (list alph 0 (list) refB))
(define rC (list alph 0 (list) refC))
(define a (list alph 0 (list) alph))

(define (make-rotor rotor times-rotated notch)
  (list rotor times-rotated notch))
(define (get-rotor-alph r) (car r))
(define (get-times-rotated r) (cadr r))
(define (get-rotor-notch r) (caddr r))
(define (get-rotor-encrypt r) (cadddr r))
;; when rotor rotates: the times increments by 1, the rotor rotates once
(define (rotate r) (list (rotate-rotor 1 (get-rotor-alph r))
                         (modulo (+ 1 (get-times-rotated r)) 26)
                         (get-rotor-notch r)
                         (get-rotor-encrypt r)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; rotating the rotor
(define (rotate-rotor times rotor)
  ;; shift it all up one
  (define (iter times to-shift r )
    (if (= 0 times) (append r to-shift)
        (iter (- times 1) (append to-shift (list (car r))) (cdr r))))
  (iter times (list) rotor))

;; the letter passed to the next rotor will not fall onto the same position
;; it depends on the number of times this rotor rotated and the next rotor
;; rotated; encryption before the reflector

(define (get-idx-next-letter-forward r1 r2 idx-letter-r1)
  (get-alph-idx (modulo (+ (- idx-letter-r1 (get-times-rotated r1))
             26
             (get-times-rotated r2))
          26)))
;; encryption after the reflector
(define (get-idx-next-letter-backward r2 r1 idx-letter-r2)
  (get-alph-idx (modulo (- (+
              idx-letter-r2
              (get-times-rotated r1))
             (get-times-rotated r2))
          26)))
;; the 0th (or 26th) index becomes Z, the first index is A
(define (get-alph-idx idx-letter)
  (if (= 0 idx-letter) 26 idx-letter))

;; letter must be a char. Alphabet is a list from a string (list of chars)
;; gets the index of a specific letter in a sequence of letters
(define (get-idx-of letter alphabet)
  (define (iter count a)
    (if (empty? a) -1
    (if (equal? (car a) letter)
        count
        (iter (+ 1 count) (cdr a)))))
  (iter 0 alphabet))

;; gets letter at a specific index in a sequence of letters
(define (get idx alphabet)
  (define (iter count a)
    (if (= 0 count) (car a) (iter (- count 1) (cdr a))))
  (iter idx alphabet))

(define (pass-through-first-rotor r letter)
  ;; returns the index + 1 of the letter in the alphabet to be passed to
  ;; the second rotor
  (+ 1 (get-idx-of
        (get (get-idx-of
              (get (get-idx-of letter alph) (get-rotor-alph r))
              alph)
             (get-rotor-encrypt r))
        alph))
  )

;; returns the index of the letter starting from 1 (to use in index searching,
;; must subtract 1)
;; if r2 is empty, the idx-r2 should be from 0, otherwise, it is from 1
;; r2 is never going to be empty because if it is the last rotor, the reflector
;; will be passed
(define (pass-through-rotor-forward r1 idx-r2 . r2)
  (let ((idx (if (empty? r2) idx-r2
                 (- (get-idx-next-letter-forward (car r2) r1 idx-r2) 1))))
      (+ 1 (get-idx-of
        (get (get-idx-of
              (get idx (get-rotor-alph r1))
              alph)
             (get-rotor-encrypt r1))
        alph))))

;; index is from 1; returns index from 1
(define (pass-through-rotor-backward r1 idx-r2 . r2)
  (let ((idx (- (get-idx-next-letter-backward (car r2) r1 idx-r2) 1)))
    (+ 1(get-idx-of (get
                      (get-idx-of (get idx alph)
                                  (get-rotor-encrypt r1))
                      (get-rotor-alph r1))
                (get-rotor-alph r1)))))

;; returns idx from 0
(define (pass-forward r1 r2 r3 ref letter)
  (define idx (get-idx-of letter alph))
 (- (get-idx-next-letter-forward r3 ref (pass-through-rotor-forward
                      r3
                      (pass-through-rotor-forward
                       r2
                       (pass-through-rotor-forward r1 idx)
                       r1)
                      r2)) 1))
;; idx is from 0
(define (pass-through-ref ref idx)
  (+ 1 (get-idx-of (get idx (get-rotor-encrypt ref)) alph)))

;; idx is from 1 return letter
(define (pass-backward r1 r2 r3 ref idx)
  (get (- (get-idx-next-letter-backward r1 a (pass-through-rotor-backward
                   r1
                   (pass-through-rotor-backward
                    r2
                    (pass-through-rotor-backward
                     r3 idx ref)
                    r3)
                   r2)) 1) alph))

(define (encrypt-letter r1 r2 r3 ref letter)
  (pass-backward r1 r2 r3 ref
                 (pass-through-ref ref (pass-forward r1 r2 r3 ref letter))))

(define (encrypt r1 r2 r3 ref string)
  (define str (string->list string))
  (define (iter r1 r2 r3 ref string encoded)
    ;; first must rotate all the rotors
    (if (empty? string) (list->string encoded)
        (if (= -1 (get-idx-of (car string) alph))
            (iter r1 r2 r3 ref (cdr string) (append encoded
                                                    (string->list " ")))
            (let* ((r1first (car (get-rotor-alph r1)))
               (r2first (car (get-rotor-alph r2)))
               (r11 (rotate r1))
               (r21 (if (or (equal? r1first (get-rotor-notch r1))
                            (equal? r2first (get-rotor-notch r2)))
                        (rotate r2) r2))
               (r31 (if (equal? r2first (get-rotor-notch r2))
                        (rotate r3) r3)))
          (iter r11 r21 r31 ref (cdr string)
                (append encoded (list (encrypt-letter r11 r21 r31 ref
                                                      (car string)))))))))
  (iter r1 r2 r3 ref str (list)))

(provide alph rI rII rIII rIV rV rB rC get-idx-of get
         get-idx-next-letter-forward get-idx-next-letter-backward
         rotate pass-through-first-rotor pass-through-rotor-forward
         pass-through-rotor-backward encrypt-letter pass-forward
         pass-through-ref pass-backward encrypt)