diff --git a/Enigma/Rotor.swift b/Enigma/Rotor.swift index fa6c2a7..1f0d83f 100644 --- a/Enigma/Rotor.swift +++ b/Enigma/Rotor.swift @@ -9,18 +9,32 @@ import Foundation -class Rotor { +protocol Encoder { + func encode(c: Character) throws -> Character +} + + +class Cryptor { + enum Error: ErrorType { + /** Thrown when encode() encounters a character that is not in the alphabet. */ + case InvalidCharacter(c: Character) + } + + /** Array of all possible characters to encrypt. */ + static let alphabet: [Character] = Array("ABCDEFGHIJKLMNOPQRSTUVWXYZ".characters) +} + + +/** Rotors are Cryptors that have a position, which offsets the alphabet from the series and changes which character is substituted for a given input. */ +class Rotor: Cryptor, Encoder { enum Error: ErrorType { /** Thrown when the initializer is given an invalid series. */ case InvalidSeries - /** Thrown when encode() encounters a character that is not in the alphabet. */ - case InvalidCharacter } - static let alphabet: [Character] = Array("ABCDEFGHIJKLMNOPQRSTUVWXYZ".characters) - /** The position of first letter in `series` in the `alphabet`. */ - var position: Int + var position: Int = 0 + /** The series of characters that this rotor cycles through. */ let series: [Character]! @@ -29,12 +43,11 @@ class Rotor { } init(series: [Character]) throws { - self.position = 0 - guard series.count == Rotor.alphabet.count else { - self.series = nil + self.series = series + super.init() + guard series.count == Cryptor.alphabet.count else { throw Error.InvalidSeries } - self.series = series } func advance(count: Int = 1) { @@ -44,8 +57,64 @@ class Rotor { func encode(c: Character) throws -> Character { let offset: Int! = Rotor.alphabet.indexOf(c) guard offset != nil else { - throw Error.InvalidCharacter + throw Error.InvalidCharacter(c: c) } return series[(offset + position) % series.count] } +} + + +/** A Plugboard is a Cryptor that substitutes one character for another based on a set of pairs. A pair of characters is mutually exclusive of other pairs; that is, a character can only belong to one pair. Furthermore, the Plugboard always trades one character for the same character and vice versa. */ +class Plugboard: Cryptor, Encoder { + enum Error: ErrorType { + case CharacterInMultiplePairs(c: Character) + case PairContainsSameCharacter(c: Character) + } + + var pairs: [(Character, Character)]! + + init(pairs: [(Character, Character)] = []) throws { + super.init() + try validatePairs(pairs) + self.pairs = pairs + } + + func validatePairs(pairs: [(Character, Character)]) throws { + let alphabetSet = Set(Cryptor.alphabet) + var charactersSeen = Set() + for pair in pairs { + if !alphabetSet.contains(pair.0) { + throw Cryptor.Error.InvalidCharacter(c: pair.0) + } + if !alphabetSet.contains(pair.1) { + throw Cryptor.Error.InvalidCharacter(c: pair.1) + } + if pair.0 == pair.1 { + throw Error.PairContainsSameCharacter(c: pair.0) + } + if charactersSeen.contains(pair.0) { + throw Error.CharacterInMultiplePairs(c: pair.0) + } else { + charactersSeen.insert(pair.0) + } + if charactersSeen.contains(pair.1) { + throw Error.CharacterInMultiplePairs(c: pair.1) + } else { + charactersSeen.insert(pair.1) + } + } + } + + func encode(c: Character) throws -> Character { + for pair in pairs { + if c == pair.0 { + return pair.1 + } + if c == pair.1 { + return pair.0 + } + } + // If no pair exists, just return the character itself. + return c + } } \ No newline at end of file diff --git a/EnigmaTests/EnigmaTests.swift b/EnigmaTests/EnigmaTests.swift index b34248d..faf5c27 100644 --- a/EnigmaTests/EnigmaTests.swift +++ b/EnigmaTests/EnigmaTests.swift @@ -9,8 +9,9 @@ import XCTest @testable import Enigma +let alphaSeries = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + class RotorTests: XCTestCase { - let alphaSeries = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" let rotorSeries = "EKMFLGDQVZNTOWYHXUSPAIBRCJ" func testThatUnadvancedSubstitutionWorks() { @@ -28,4 +29,20 @@ class RotorTests: XCTestCase { XCTAssertEqual(try! rotor.encode(plainCharacter), cipherCharacter) } } +} + + +class PlugboardTests: XCTestCase { + func testThatEmptyPlugboardPassesThroughAllCharacters() { + let plugboard = try! Enigma.Plugboard() + for c in alphaSeries.characters { + XCTAssertEqual(try! plugboard.encode(c), c) + } + } + + func testThatPlugboardPairsAreBidirectional() { + let plugboard = try! Enigma.Plugboard(pairs: [("A", "H")]) + XCTAssertEqual(try! plugboard.encode("A"), "H") + XCTAssertEqual(try! plugboard.encode("H"), "A") + } } \ No newline at end of file