enigma.swift

import Foundation

enum Token: Character {
    case A = "A"
    case B = "B"
    case C = "C"
    case D = "D"
    case E = "E"
    case F = "F"
    case G = "G"
    case H = "H"
    case I = "I"
    case J = "J"
    case K = "K"
    case L = "L"
    case M = "M"
    case N = "N"
    case O = "O"
    case P = "P"
    case Q = "Q"
    case R = "R"
    case S = "S"
    case T = "T"
    case U = "U"
    case V = "V"
    case W = "W"
    case X = "X"
    case Y = "Y"
    case Z = "Z"
}

extension Token {
    func next() -> Token {
        switch self {
        case .A: return .B
        case .B: return .C
        case .C: return .D
        case .D: return .E
        case .E: return .F
        case .F: return .G
        case .G: return .H
        case .H: return .I
        case .I: return .J
        case .J: return .K
        case .K: return .L
        case .L: return .M
        case .M: return .N
        case .N: return .O
        case .O: return .P
        case .P: return .Q
        case .Q: return .R
        case .R: return .S
        case .S: return .T
        case .T: return .U
        case .U: return .V
        case .V: return .W
        case .W: return .X
        case .X: return .Y
        case .Y: return .Z
        case .Z: return .A
        }
    }
}
extension Token {
    var num: Int {
        switch self {
        case .A: return 0
        case .B: return 1
        case .C: return 2
        case .D: return 3
        case .E: return 4
        case .F: return 5
        case .G: return 6
        case .H: return 7
        case .I: return 8
        case .J: return 9
        case .K: return 10
        case .L: return 11
        case .M: return 12
        case .N: return 13
        case .O: return 14
        case .P: return 15
        case .Q: return 16
        case .R: return 17
        case .S: return 18
        case .T: return 19
        case .U: return 20
        case .V: return 21
        case .W: return 22
        case .X: return 23
        case .Y: return 24
        case .Z: return 25
        }
    }
}
extension Int {
    var token: Token {
        switch self {
        case 0: return .A
        case 1: return .B
        case 2: return .C
        case 3: return .D
        case 4: return .E
        case 5: return .F
        case 6: return .G
        case 7: return .H
        case 8: return .I
        case 9: return .J
        case 10: return .K
        case 11: return .L
        case 12: return .M
        case 13: return .N
        case 14: return .O
        case 15: return .P
        case 16: return .Q
        case 17: return .R
        case 18: return .S
        case 19: return .T
        case 20: return .U
        case 21: return .V
        case 22: return .W
        case 23: return .X
        case 24: return .Y
        case 25: return .Z
        default:
            fatalError()
        }
    }
}

extension Token {
    static func - (lhs: Token, rhs: Int) -> Token {
        return ((lhs.num - rhs + 26) % 26).token
    }
    static func - (lhs: Token, rhs: Token) -> Int {
        return lhs.num - rhs.num
    }
}

// -------------------------------------------- //

typealias Cipher = (Token) -> (Token)

struct RotorBox {
    
    struct Rotor {
        
        let forward: Cipher
        let backward: Cipher
        var position: Token
        
        init(_ forward: @escaping Cipher, position: Token) {
            self.forward = forward
            self.backward = { token in
                let tokens: [Token] = [.A, .B, .C, .D, .E, .F, .G, .H, .I, .J, .K, .L, .M, .N, .O, .P, .Q, .R, .S, .T, .U, .V, .W, .X, .Y, .Z]
                return tokens.filter{ forward($0) == token }.first!
            }
            self.position = position
        }
        
        mutating func rotate() {
            position = position.next()
        }
    }
    
    private var pure = Rotor({ $0 }, position: .A)
    private var rotor0: Rotor
    private var rotor1: Rotor
    private var rotor2: Rotor
    private let reflector = Rotor({ token in
        switch token {
        case .A: return .E
        case .B: return .J
        case .C: return .M
        case .D: return .Z
        case .E: return .A
        case .F: return .L
        case .G: return .Y
        case .H: return .X
        case .I: return .V
        case .J: return .B
        case .K: return .W
        case .L: return .F
        case .M: return .C
        case .N: return .R
        case .O: return .Q
        case .P: return .U
        case .Q: return .O
        case .R: return .N
        case .S: return .T
        case .T: return .S
        case .U: return .P
        case .V: return .I
        case .W: return .K
        case .X: return .H
        case .Y: return .G
        case .Z: return .D
        }
    }, position: .A) // regard the position of the reflector as .A
    
    private var mod1 = 0
    private var mod2 = 0
    
    init(rotor0: Rotor, rotor1: Rotor, rotor2: Rotor) {
        self.rotor0 = rotor0
        self.rotor1 = rotor1
        self.rotor2 = rotor2
    }
    
    mutating func cipher(_ target: Token) -> Token {
        defer {
            rotate()
        }
        
        let input = connect(pure, rotor0)(target)
        let r0 = rotor0.forward(input)
        let _r0 = connect(rotor0, rotor1)(r0)
        let r1 = rotor1.forward(_r0)
        let _r1 = connect(rotor1, rotor2)(r1)
        let r2 = rotor2.forward(_r1)
        let _r2 = connect(rotor2, reflector)(r2)
        let ref = reflector.forward(_r2)
        let _ref = connect(reflector, rotor2)(ref)
        let r2ret = rotor2.backward(_ref)
        let _r2ret = connect(rotor2, rotor1)(r2ret)
        let r1ret = rotor1.backward(_r2ret)
        let _r1ret = connect(rotor1, rotor0)(r1ret)
        let r0ret = rotor0.backward(_r1ret)
        let output = connect(rotor0, pure)(r0ret)
        
        return output
    }
    
    mutating func rotate() {
        mod1 = (mod1 + 1) % 26
        mod2 = (mod2 + 1) % (26 * 26)
        rotor0.rotate()
        if mod1 == 0 {
            rotor1.rotate()
        }
        if mod2 == 0 {
            rotor2.rotate()
        }
    }
    
    func connect(_ input: Rotor, _ output: Rotor) -> (Token) -> Token {
        let gap = input.position - output.position
        return { token in
            return token - gap
        }
    }
}

struct Enigma {
    
    private var rotorBox: RotorBox
    private let plugboard: Cipher
    
    typealias Key = (Token, Token, Token)
    init(rotor0: @escaping Cipher, rotor1: @escaping Cipher, rotor2: @escaping Cipher, plugboard: @escaping Cipher, key: Key) {
        self.rotorBox = RotorBox(
            rotor0: RotorBox.Rotor(rotor0, position: key.0),
            rotor1: RotorBox.Rotor(rotor1, position: key.1),
            rotor2: RotorBox.Rotor(rotor2, position: key.2)
        )
        self.plugboard = plugboard
    }
    
    mutating func cipher(_ target: Token) -> Token {
        return plugboard(rotorBox.cipher(plugboard(target)))
    }
    
    mutating func cipher(_ target: String) -> String {
        var output = ""
        for char in target.characters {
            output += String(cipher(Token(rawValue: char)!).rawValue)
        }
        return output
    }
}


// ----------------------------------------------------- //

// https://en.wikipedia.org/wiki/Enigma_rotor_details#Rotor_wiring_tables
// Swiss K models
// Reflector A

let I_K: Cipher = { token in
    switch token {
    case .A: return .P
    case .B: return .E
    case .C: return .Z
    case .D: return .U
    case .E: return .O
    case .F: return .H
    case .G: return .X
    case .H: return .S
    case .I: return .C
    case .J: return .V
    case .K: return .F
    case .L: return .M
    case .M: return .T
    case .N: return .B
    case .O: return .G
    case .P: return .L
    case .Q: return .R
    case .R: return .I
    case .S: return .N
    case .T: return .Q
    case .U: return .J
    case .V: return .W
    case .W: return .A
    case .X: return .Y
    case .Y: return .D
    case .Z: return .K
    }
}
let II_K: Cipher = { token in
    switch token {
    case .A: return .Z
    case .B: return .O
    case .C: return .U
    case .D: return .E
    case .E: return .S
    case .F: return .Y
    case .G: return .D
    case .H: return .K
    case .I: return .F
    case .J: return .W
    case .K: return .P
    case .L: return .I
    case .M: return .C
    case .N: return .Q
    case .O: return .X
    case .P: return .H
    case .Q: return .M
    case .R: return .V
    case .S: return .B
    case .T: return .L
    case .U: return .G
    case .V: return .N
    case .W: return .J
    case .X: return .R
    case .Y: return .A
    case .Z: return .T
    }
}
let III_K: Cipher = { token in
    switch token {
    case .A: return .E
    case .B: return .H
    case .C: return .R
    case .D: return .V
    case .E: return .X
    case .F: return .G
    case .G: return .A
    case .H: return .O
    case .I: return .B
    case .J: return .Q
    case .K: return .U
    case .L: return .S
    case .M: return .I
    case .N: return .M
    case .O: return .Z
    case .P: return .F
    case .Q: return .L
    case .R: return .Y
    case .S: return .N
    case .T: return .W
    case .U: return .K
    case .V: return .T
    case .W: return .P
    case .X: return .D
    case .Y: return .J
    case .Z: return .C
    }
}
let UKW_K: Cipher = { token in
    switch token {
    case .A: return .I
    case .B: return .M
    case .C: return .E
    case .D: return .T
    case .E: return .C
    case .F: return .G
    case .G: return .F
    case .H: return .R
    case .I: return .A
    case .J: return .Y
    case .K: return .S
    case .L: return .Q
    case .M: return .B
    case .N: return .Z
    case .O: return .X
    case .P: return .W
    case .Q: return .L
    case .R: return .H
    case .S: return .K
    case .T: return .D
    case .U: return .V
    case .V: return .U
    case .W: return .P
    case .X: return .O
    case .Y: return .J
    case .Z: return .N
    }
}
let ETW_K: Cipher = { token in
    switch token {
    case .A: return .Q
    case .B: return .W
    case .C: return .E
    case .D: return .R
    case .E: return .T
    case .F: return .Z
    case .G: return .U
    case .H: return .I
    case .I: return .O
    case .J: return .A
    case .K: return .S
    case .L: return .D
    case .M: return .F
    case .N: return .G
    case .O: return .H
    case .P: return .J
    case .Q: return .K
    case .R: return .P
    case .S: return .Y
    case .T: return .X
    case .U: return .C
    case .V: return .V
    case .W: return .B
    case .X: return .N
    case .Y: return .M
    case .Z: return .L
    }
}
let swissK = [I_K, II_K, III_K, UKW_K, ETW_K]
let plugboard: Cipher = { token in
    switch token {
    default:
        return token
    }
}

var enigma = Enigma(rotor0: swissK[0], rotor1: swissK[1], rotor2: swissK[2], plugboard: plugboard, key: (.A, .A, .A))
let message = "HELLOWORLD"
let ciphered = enigma.cipher(message)

var _enigma = Enigma(rotor0: swissK[0], rotor1: swissK[1], rotor2: swissK[2], plugboard: plugboard, key: (.A, .A, .A))
let deciphered = _enigma.cipher(ciphered)