{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Enigma machine\n",
    "Specification from [Codes and Ciphers](http://www.codesandciphers.org.uk/enigma/rotorspec.htm) page.\n",
    "\n",
    "Example Enigma machines from [Louise Dale](http://enigma.louisedade.co.uk/enigma.html) (full simulation) and [EnigmaCo](http://enigmaco.de/enigma/enigma.html) (good animation of the wheels, but no ring settings).\n",
    "\n",
    "There's also the nice Enigma simulator for Android by [Franklin Heath](https://franklinheath.co.uk/2012/02/04/our-first-app-published-enigma-simulator/), available on the [Google Play store](https://play.google.com/store/apps/details?id=uk.co.franklinheath.enigmasim&hl=en_GB)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 333,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import string\n",
    "\n",
    "cat = ''.join"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 334,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "wheel_i_spec = 'ekmflgdqvzntowyhxuspaibrcj'\n",
    "wheel_ii_spec = 'ajdksiruxblhwtmcqgznpyfvoe'\n",
    "wheel_iii_spec = 'bdfhjlcprtxvznyeiwgakmusqo'\n",
    "wheel_iv_spec = 'esovpzjayquirhxlnftgkdcmwb'\n",
    "wheel_v_spec = 'vzbrgityupsdnhlxawmjqofeck'\n",
    "wheel_vi_spec = 'jpgvoumfyqbenhzrdkasxlictw'\n",
    "wheel_vii_spec = 'nzjhgrcxmyswboufaivlpekqdt'\n",
    "wheel_viii_spec = 'fkqhtlxocbjspdzramewniuygv'\n",
    "beta_wheel_spec = 'leyjvcnixwpbqmdrtakzgfuhos'\n",
    "gamma_wheel_spec = 'fsokanuerhmbtiycwlqpzxvgjd'\n",
    "\n",
    "wheel_i_pegs = ['q']\n",
    "wheel_ii_pegs = ['e']\n",
    "wheel_iii_pegs = ['v']\n",
    "wheel_iv_pegs = ['j']\n",
    "wheel_v_pegs = ['z']\n",
    "wheel_vi_pegs = ['z', 'm']\n",
    "wheel_vii_pegs = ['z', 'm']\n",
    "wheel_viii_pegs = ['z', 'm']\n",
    "\n",
    "reflector_b_spec = '(ay) (br) (cu) (dh) (eq) (fs) (gl) (ip) (jx) (kn) (mo) (tz) (vw)'\n",
    "reflector_c_spec = '(af) (bv) (cp) (dj) (ei) (go) (hy) (kr) (lz) (mx) (nw) (tq) (su)'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 335,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class LetterTransformer(object):\n",
    "    def __init__(self, specification, raw_transform=False):\n",
    "        if raw_transform:\n",
    "            transform = specification\n",
    "        else:\n",
    "            transform = self.parse_specification(specification)\n",
    "        self.validate_transform(transform)\n",
    "        self.make_transform_map(transform)\n",
    "    \n",
    "    def parse_specification(self, specification):\n",
    "        return specification\n",
    "    \n",
    "    def validate_transform(self, transform):\n",
    "        \"\"\"A set of pairs, of from-to\"\"\"\n",
    "        if len(transform) != 26:\n",
    "            raise ValueError(\"Transform specification has {} pairs, requires 26\".\n",
    "                format(len(transform)))\n",
    "        for p in transform:\n",
    "            if len(p) != 2:\n",
    "                raise ValueError(\"Not all mappings in transform \"\n",
    "                    \"have two elements\")\n",
    "        if len(set([p[0] for p in transform])) != 26:\n",
    "            raise ValueError(\"Transform specification must list 26 origin letters\") \n",
    "        if len(set([p[1] for p in transform])) != 26:\n",
    "            raise ValueError(\"Transform specification must list 26 destination letters\") \n",
    "\n",
    "    def make_empty_transform(self):\n",
    "        self.forward_map = [0] * 26\n",
    "        self.backward_map = [0] * 26\n",
    "            \n",
    "    def make_transform_map(self, transform):\n",
    "        self.make_empty_transform()\n",
    "        for p in transform:\n",
    "            self.forward_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')\n",
    "            self.backward_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')\n",
    "        return self.forward_map, self.backward_map\n",
    "    \n",
    "    def forward(self, letter):\n",
    "        if letter in string.ascii_lowercase:\n",
    "            return chr(\n",
    "                (self.forward_map[(ord(letter) - ord('a')) % 26] + ord('a')))\n",
    "        else:\n",
    "            return ''\n",
    "                \n",
    "    def backward(self, letter):\n",
    "        if letter in string.ascii_lowercase:\n",
    "            return chr(\n",
    "                (self.backward_map[(ord(letter) - ord('a')) % 26] + ord('a')))\n",
    "        else:\n",
    "            return ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 336,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('z', 'a'),\n",
       " ('a', 'b'),\n",
       " ('b', 'c'),\n",
       " ('c', 'd'),\n",
       " ('d', 'e'),\n",
       " ('e', 'f'),\n",
       " ('f', 'g'),\n",
       " ('g', 'h'),\n",
       " ('h', 'i'),\n",
       " ('i', 'j'),\n",
       " ('j', 'k'),\n",
       " ('k', 'l'),\n",
       " ('l', 'm'),\n",
       " ('m', 'n'),\n",
       " ('n', 'o'),\n",
       " ('o', 'p'),\n",
       " ('p', 'q'),\n",
       " ('q', 'r'),\n",
       " ('r', 's'),\n",
       " ('s', 't'),\n",
       " ('t', 'u'),\n",
       " ('u', 'v'),\n",
       " ('v', 'w'),\n",
       " ('w', 'x'),\n",
       " ('x', 'y'),\n",
       " ('y', 'z')]"
      ]
     },
     "execution_count": 336,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tmap = [('z', 'a')] + [(l, string.ascii_lowercase[i+1]) for i, l in enumerate(string.ascii_lowercase[:-1])]\n",
    "tmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 337,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([1,\n",
       "  2,\n",
       "  3,\n",
       "  4,\n",
       "  5,\n",
       "  6,\n",
       "  7,\n",
       "  8,\n",
       "  9,\n",
       "  10,\n",
       "  11,\n",
       "  12,\n",
       "  13,\n",
       "  14,\n",
       "  15,\n",
       "  16,\n",
       "  17,\n",
       "  18,\n",
       "  19,\n",
       "  20,\n",
       "  21,\n",
       "  22,\n",
       "  23,\n",
       "  24,\n",
       "  25,\n",
       "  0],\n",
       " [25,\n",
       "  0,\n",
       "  1,\n",
       "  2,\n",
       "  3,\n",
       "  4,\n",
       "  5,\n",
       "  6,\n",
       "  7,\n",
       "  8,\n",
       "  9,\n",
       "  10,\n",
       "  11,\n",
       "  12,\n",
       "  13,\n",
       "  14,\n",
       "  15,\n",
       "  16,\n",
       "  17,\n",
       "  18,\n",
       "  19,\n",
       "  20,\n",
       "  21,\n",
       "  22,\n",
       "  23,\n",
       "  24])"
      ]
     },
     "execution_count": 337,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lt = LetterTransformer(tmap)\n",
    "lt.forward_map, lt.backward_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 338,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'bcdefghijklmnopqrstuvwxyza'"
      ]
     },
     "execution_count": 338,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(lt.forward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 339,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'zabcdefghijklmnopqrstuvwxy'"
      ]
     },
     "execution_count": 339,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(lt.backward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 340,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Plugboard(LetterTransformer):\n",
    "    def parse_specification(self, specification):\n",
    "        return [tuple(p) for p in specification.split()]\n",
    "    \n",
    "    def validate_transform(self, transform):\n",
    "        \"\"\"A set of pairs, of from-to\"\"\"\n",
    "        for p in transform:\n",
    "            if len(p) != 2:\n",
    "                raise ValueError(\"Not all mappings in transform\"\n",
    "                    \"have two elements\")\n",
    "    \n",
    "    def make_empty_transform(self):\n",
    "        self.forward_map = list(range(26))\n",
    "        self.backward_map = list(range(26))\n",
    "        \n",
    "    def make_transform_map(self, transform):\n",
    "        expanded_transform = transform + [tuple(reversed(p)) for p in transform]\n",
    "        return super(Plugboard, self).make_transform_map(expanded_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 341,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pb = Plugboard([('a', 'z'), ('b', 'y')], raw_transform=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 342,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'zycdefghijklmnopqrstuvwxba'"
      ]
     },
     "execution_count": 342,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(pb.forward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'zycdefghijklmnopqrstuvwxba'"
      ]
     },
     "execution_count": 343,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(pb.backward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pb = Plugboard('az by')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('zycdefghijklmnopqrstuvwxba', 'zycdefghijklmnopqrstuvwxba')"
      ]
     },
     "execution_count": 345,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(pb.forward(l) for l in string.ascii_lowercase), cat(pb.backward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 346,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Reflector(Plugboard):\n",
    "    def parse_specification(self, specification):\n",
    "        return [tuple(p) for p in specification[1:-1].split(') (')]\n",
    "    \n",
    "    def validate_transform(self, transform):\n",
    "        if len(transform) != 13:\n",
    "            raise ValueError(\"Reflector specification has {} pairs, requires 13\".\n",
    "                format(len(transform)))\n",
    "        if len(set([p[0] for p in transform] + \n",
    "                    [p[1] for p in transform])) != 26:\n",
    "            raise ValueError(\"Reflector specification does not contain 26 letters\")\n",
    "        try:\n",
    "            super(Reflector, self).validate_transform(transform)\n",
    "        except ValueError as v:\n",
    "            raise ValueError(\"Not all mappings in reflector have two elements\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 347,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('a', 'y'),\n",
       " ('b', 'r'),\n",
       " ('c', 'u'),\n",
       " ('d', 'h'),\n",
       " ('e', 'q'),\n",
       " ('f', 's'),\n",
       " ('g', 'l'),\n",
       " ('i', 'p'),\n",
       " ('j', 'x'),\n",
       " ('k', 'n'),\n",
       " ('m', 'o'),\n",
       " ('t', 'z'),\n",
       " ('v', 'w')]"
      ]
     },
     "execution_count": 347,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# reflector_b_text = '(AY) (BR) (CU) (DH) (EQ) (FS) (GL) (IP) (JX) (KN) (MO) (TZ) (VW)'\n",
    "reflector_b_l = [tuple(p) for p in reflector_b_spec.lower()[1:-1].split(') (')]\n",
    "reflector_b_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 348,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "reflector_b = Reflector(reflector_b_spec, raw_transform=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 349,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'yruhqsldpxngokmiebfzcwvjat'"
      ]
     },
     "execution_count": 349,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(reflector_b.forward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 350,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "reflector_c = Reflector(reflector_c_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 351,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'fvpjiaoyedrzxwgctkuqsbnmhl'"
      ]
     },
     "execution_count": 351,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(reflector_c.forward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 352,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class SimpleWheel(LetterTransformer):\n",
    "    def __init__(self, transform, position='a', raw_transform=False):\n",
    "        super(SimpleWheel, self).__init__(transform, raw_transform)\n",
    "        self.set_position(position)\n",
    "        \n",
    "    def parse_specification(self, specification):\n",
    "        return list(zip(string.ascii_lowercase, specification))\n",
    "    \n",
    "    def set_position(self, position):\n",
    "        self.position = ord(position) - ord('a')\n",
    "        self.position_l = position\n",
    "    \n",
    "    def forward(self, letter):\n",
    "        if letter in string.ascii_lowercase:\n",
    "            return chr(\n",
    "                (self.forward_map[(ord(letter) - ord('a') + self.position) % 26] - \n",
    "                    self.position) % 26 + \n",
    "                ord('a'))\n",
    "        else:\n",
    "            return ''\n",
    "                \n",
    "    def backward(self, letter):\n",
    "        if letter in string.ascii_lowercase:\n",
    "            return chr(\n",
    "                (self.backward_map[(ord(letter) - ord('a') + self.position) % 26] - \n",
    "                    self.position) % 26 + \n",
    "                ord('a'))\n",
    "        else:\n",
    "            return ''\n",
    "        \n",
    "    def advance(self):\n",
    "        self.position = (self.position + 1) % 26\n",
    "        self.position_l = chr(self.position + ord('a'))\n",
    "        return self.position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 353,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('a', 'e'),\n",
       " ('b', 'k'),\n",
       " ('c', 'm'),\n",
       " ('d', 'f'),\n",
       " ('e', 'l'),\n",
       " ('f', 'g'),\n",
       " ('g', 'd'),\n",
       " ('h', 'q'),\n",
       " ('i', 'v'),\n",
       " ('j', 'z'),\n",
       " ('k', 'n'),\n",
       " ('l', 't'),\n",
       " ('m', 'o'),\n",
       " ('n', 'w'),\n",
       " ('o', 'y'),\n",
       " ('p', 'h'),\n",
       " ('q', 'x'),\n",
       " ('r', 'u'),\n",
       " ('s', 's'),\n",
       " ('t', 'p'),\n",
       " ('u', 'a'),\n",
       " ('v', 'i'),\n",
       " ('w', 'b'),\n",
       " ('x', 'r'),\n",
       " ('y', 'c'),\n",
       " ('z', 'j')]"
      ]
     },
     "execution_count": 353,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rotor_1_transform = list(zip(string.ascii_lowercase, 'EKMFLGDQVZNTOWYHXUSPAIBRCJ'.lower()))\n",
    "rotor_1_transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 354,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "wheel_1 = SimpleWheel(rotor_1_transform, raw_transform=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 355,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('ekmflgdqvzntowyhxuspaibrcj', 'uwygadfpvzbeckmthxslrinqoj')"
      ]
     },
     "execution_count": 355,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(wheel_1.forward(l) for l in string.ascii_lowercase), cat(wheel_1.backward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 356,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "wheel_2 = SimpleWheel(wheel_ii_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 357,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('ajdksiruxblhwtmcqgznpyfvoe', 'ajpczwrlfbdkotyuqgenhxmivs')"
      ]
     },
     "execution_count": 357,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat(wheel_2.forward(l) for l in string.ascii_lowercase), cat(wheel_2.backward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 358,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cegikboqswuymxdhvfzjltrpna'"
      ]
     },
     "execution_count": 358,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wheel_3 = SimpleWheel(wheel_iii_spec)\n",
    "wheel_3.set_position('a')\n",
    "wheel_3.advance()\n",
    "cat(wheel_3.forward(l) for l in string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 524,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Wheel(SimpleWheel):\n",
    "    def __init__(self, transform, ring_peg_letters, ring_setting=1, position='a', raw_transform=False):\n",
    "        self.ring_peg_letters = ring_peg_letters\n",
    "        self.ring_setting = ring_setting\n",
    "        super(Wheel, self).__init__(transform, position=position, raw_transform=raw_transform)\n",
    "        self.set_position(position)\n",
    "        \n",
    "    def set_position(self, position):\n",
    "        super(Wheel, self).set_position(position)\n",
    "        self.peg_positions = [(ord(p) - ord(position)) % 26  for p in self.ring_peg_letters]\n",
    "        self.position = (self.position - self.ring_setting + 1) % 26\n",
    "        \n",
    "    def advance(self):\n",
    "        super(Wheel, self).advance()\n",
    "        self.peg_positions = [(p - 1) % 26 for p in self.peg_positions]\n",
    "        return self.position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 525,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "wheel_3 = Wheel(wheel_iii_spec, wheel_iii_pegs, position='b', ring_setting=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 526,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, [20])"
      ]
     },
     "execution_count": 526,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wheel_3.position, wheel_3.peg_positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 527,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(25, [24, 11])"
      ]
     },
     "execution_count": 527,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wheel_6 = Wheel(wheel_vi_spec, wheel_vi_pegs, position='b', ring_setting=3)\n",
    "wheel_6.position, wheel_6.peg_positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 528,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 [23, 10]\n",
      "1 [22, 9]\n",
      "2 [21, 8]\n",
      "3 [20, 7]\n",
      "4 [19, 6]\n",
      "5 [18, 5]\n",
      "6 [17, 4]\n",
      "7 [16, 3]\n",
      "8 [15, 2]\n",
      "9 [14, 1]\n",
      "10 [13, 0]\n",
      "11 [12, 25]\n",
      "12 [11, 24]\n",
      "13 [10, 23]\n",
      "14 [9, 22]\n",
      "15 [8, 21]\n",
      "16 [7, 20]\n",
      "17 [6, 19]\n",
      "18 [5, 18]\n",
      "19 [4, 17]\n",
      "20 [3, 16]\n",
      "21 [2, 15]\n",
      "22 [1, 14]\n",
      "23 [0, 13]\n",
      "24 [25, 12]\n",
      "25 [24, 11]\n",
      "0 [23, 10]\n"
     ]
    }
   ],
   "source": [
    "for _ in range(27):\n",
    "    wheel_6.advance()\n",
    "    print(wheel_6.position, wheel_6.peg_positions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 529,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Enigma(object):\n",
    "    def __init__(self, reflector_spec,\n",
    "                 left_wheel_spec, left_wheel_pegs,\n",
    "                 middle_wheel_spec, middle_wheel_pegs,\n",
    "                 right_wheel_spec, right_wheel_pegs,\n",
    "                 left_ring_setting, middle_ring_setting, right_ring_setting,\n",
    "                 plugboard_setting):\n",
    "        self.reflector = Reflector(reflector_spec)\n",
    "        self.left_wheel = Wheel(left_wheel_spec, left_wheel_pegs, ring_setting=left_ring_setting)\n",
    "        self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_pegs, ring_setting=middle_ring_setting)\n",
    "        self.right_wheel = Wheel(right_wheel_spec, right_wheel_pegs, ring_setting=right_ring_setting)\n",
    "        self.plugboard = Plugboard(plugboard_setting)\n",
    "    \n",
    "    def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):\n",
    "        self.left_wheel.set_position(left_wheel_position)\n",
    "        self.middle_wheel.set_position(middle_wheel_position)\n",
    "        self.right_wheel.set_position(right_wheel_position)\n",
    "        \n",
    "    def lookup(self, letter):\n",
    "        a = self.plugboard.forward(letter)\n",
    "        b = self.right_wheel.forward(a)\n",
    "        c = self.middle_wheel.forward(b)\n",
    "        d = self.left_wheel.forward(c)\n",
    "        e = self.reflector.forward(d)\n",
    "        f = self.left_wheel.backward(e)\n",
    "        g = self.middle_wheel.backward(f)\n",
    "        h = self.right_wheel.backward(g)\n",
    "        i = self.plugboard.backward(h)\n",
    "        return i\n",
    "    \n",
    "    def advance(self):\n",
    "        advance_middle = False\n",
    "        advance_left = False\n",
    "        if 0 in self.right_wheel.peg_positions:\n",
    "            advance_middle = True\n",
    "        if 0 in self.middle_wheel.peg_positions:\n",
    "            advance_left = True\n",
    "        if 0 in self.middle_wheel.peg_positions and 25 in self.right_wheel.peg_positions:\n",
    "            advance_middle = True\n",
    "            advance_left = True\n",
    "        self.right_wheel.advance()\n",
    "        if advance_middle: self.middle_wheel.advance()\n",
    "        if advance_left: self.left_wheel.advance()\n",
    "            \n",
    "    def encipher_letter(self, letter):\n",
    "        self.advance()\n",
    "        return self.lookup(letter)\n",
    "    \n",
    "    def encipher(self, message):\n",
    "        enciphered = ''\n",
    "        for letter in message:\n",
    "            enciphered += self.encipher_letter(letter)\n",
    "        return enciphered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 548,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "enigma = Enigma(reflector_b_spec, \n",
    "                wheel_i_spec, wheel_i_pegs,\n",
    "                wheel_ii_spec, wheel_ii_pegs,\n",
    "                wheel_iii_spec, wheel_iii_pegs,\n",
    "                1, 1, 1,\n",
    "                '')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 549,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'u'"
      ]
     },
     "execution_count": 549,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.lookup('a')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 550,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'a'"
      ]
     },
     "execution_count": 550,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.lookup('u')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 551,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a a a ; [16] [4] [21]\n",
      "a a b ; [16] [4] [20]\n",
      "a a c ; [16] [4] [19]\n",
      "a a d ; [16] [4] [18]\n",
      "a a e ; [16] [4] [17]\n",
      "a a f ; [16] [4] [16]\n",
      "a a g ; [16] [4] [15]\n",
      "a a h ; [16] [4] [14]\n",
      "a a i ; [16] [4] [13]\n",
      "a a j ; [16] [4] [12]\n",
      "a a k ; [16] [4] [11]\n",
      "a a l ; [16] [4] [10]\n",
      "a a m ; [16] [4] [9]\n",
      "a a n ; [16] [4] [8]\n",
      "a a o ; [16] [4] [7]\n",
      "a a p ; [16] [4] [6]\n",
      "a a q ; [16] [4] [5]\n",
      "a a r ; [16] [4] [4]\n",
      "a a s ; [16] [4] [3]\n",
      "a a t ; [16] [4] [2]\n",
      "a a u ; [16] [4] [1]\n",
      "a a v ; [16] [4] [0]\n",
      "a b w ; [16] [3] [25]\n",
      "a b x ; [16] [3] [24]\n",
      "a b y ; [16] [3] [23]\n",
      "a b z ; [16] [3] [22]\n"
     ]
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "for _ in range(26):\n",
    "    print(enigma.left_wheel.position_l, enigma.middle_wheel.position_l, enigma.right_wheel.position_l, ';',\n",
    "          enigma.left_wheel.peg_positions, enigma.middle_wheel.peg_positions, enigma.right_wheel.peg_positions)\n",
    "    enigma.advance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 552,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a d t ; [16] [1] [2]\n",
      "a d u ; [16] [1] [1]\n",
      "a d v ; [16] [1] [0]\n",
      "a e w ; [16] [0] [25]\n",
      "b f x ; [15] [25] [24]\n",
      "b f y ; [15] [25] [23]\n",
      "b f z ; [15] [25] [22]\n",
      "b f a ; [15] [25] [21]\n",
      "b f b ; [15] [25] [20]\n",
      "b f c ; [15] [25] [19]\n",
      "b f d ; [15] [25] [18]\n",
      "b f e ; [15] [25] [17]\n",
      "b f f ; [15] [25] [16]\n",
      "b f g ; [15] [25] [15]\n",
      "b f h ; [15] [25] [14]\n",
      "b f i ; [15] [25] [13]\n",
      "b f j ; [15] [25] [12]\n",
      "b f k ; [15] [25] [11]\n",
      "b f l ; [15] [25] [10]\n",
      "b f m ; [15] [25] [9]\n",
      "b f n ; [15] [25] [8]\n",
      "b f o ; [15] [25] [7]\n",
      "b f p ; [15] [25] [6]\n",
      "b f q ; [15] [25] [5]\n",
      "b f r ; [15] [25] [4]\n",
      "b f s ; [15] [25] [3]\n"
     ]
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'd', 't')\n",
    "for _ in range(26):\n",
    "    print(enigma.left_wheel.position_l, enigma.middle_wheel.position_l, enigma.right_wheel.position_l, ';',\n",
    "          enigma.left_wheel.peg_positions, enigma.middle_wheel.peg_positions, enigma.right_wheel.peg_positions)\n",
    "    enigma.advance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 553,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'i'"
      ]
     },
     "execution_count": 553,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "enigma.advance()\n",
    "enigma.lookup('h')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 560,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ilbdartydc'"
      ]
     },
     "execution_count": 560,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "ct = enigma.encipher('hellothere')\n",
    "assert(ct == 'ilbdartydc')\n",
    "ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 555,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'hellothere'"
      ]
     },
     "execution_count": 555,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "pt = enigma.encipher(ct)\n",
    "pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 556,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'i'"
      ]
     },
     "execution_count": 556,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "enigma.advance()\n",
    "enigma.lookup('h')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 557,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('h', 'q', 'q', 'x', 'j', 'z', 's', 'i', 'i')"
      ]
     },
     "execution_count": 557,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'a', 'a')\n",
    "enigma.advance()\n",
    "a = enigma.plugboard.forward('h')\n",
    "b = enigma.right_wheel.forward(a)\n",
    "c = enigma.middle_wheel.forward(b)\n",
    "d = enigma.left_wheel.forward(c)\n",
    "e = enigma.reflector.forward(d)\n",
    "f = enigma.left_wheel.backward(e)\n",
    "g = enigma.middle_wheel.backward(f)\n",
    "h = enigma.right_wheel.backward(g)\n",
    "i = enigma.plugboard.backward(h)\n",
    "a, b, c, d, e, f, g, h, i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 558,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'h'"
      ]
     },
     "execution_count": 558,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pb = Plugboard('')\n",
    "pb.forward('h')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 561,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'bahxvfrpdc'"
      ]
     },
     "execution_count": 561,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('a', 'd', 't')\n",
    "ct = enigma.encipher('hellothere')\n",
    "assert(ct == 'bahxvfrpdc')\n",
    "ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 562,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'kvmmwrlqlqsqpeugjrcxzwpfyiyybwloewrouvkpoztceuwtfjzqwpbqldttsr'"
      ]
     },
     "execution_count": 562,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.set_wheels('b', 'd', 'q')\n",
    "ct = enigma.encipher('aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa')\n",
    "assert(ct == 'kvmmwrlqlqsqpeugjrcxzwpfyiyybwloewrouvkpoztceuwtfjzqwpbqldttsr')\n",
    "assert(enigma.left_wheel.position_l == 'c')\n",
    "assert(enigma.middle_wheel.position_l == 'h')\n",
    "assert(enigma.right_wheel.position_l == 'a')\n",
    "ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 563,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'c'"
      ]
     },
     "execution_count": 563,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "enigma.left_wheel.position_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 564,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Setting sheet line 31 from http://www.codesandciphers.org.uk/enigma/enigma3.htm\n",
    "# Enigma simulation settings are \n",
    "# http://enigma.louisedade.co.uk/enigma.html?m3;b;b153;AFTX;AJFE;AU-BG-EY-FP-HL-IN-JZ-OS-QR-TX\n",
    "w_enigma = Enigma(reflector_b_spec, \n",
    "                wheel_i_spec, wheel_i_pegs,\n",
    "                wheel_v_spec, wheel_v_pegs,\n",
    "                wheel_iii_spec, wheel_iii_pegs,\n",
    "                6, 20, 24,\n",
    "                'ua pf rq so ni ey bg hl tx zj')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 566,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wzmfiwuntn\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'wzmfiwuntn'"
      ]
     },
     "execution_count": 566,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w_enigma.set_wheels('j', 'e', 'u')\n",
    "ct = w_enigma.encipher('hellothere')\n",
    "print(ct)\n",
    "assert(ct == 'wzmfiwuntn')\n",
    "ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 568,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ayraqvsfkhflwsmkqicvfwawswmiwvvlteb'"
      ]
     },
     "execution_count": 568,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w_enigma.set_wheels('n', 'y', 'q')\n",
    "ct = w_enigma.encipher('hellotomweshouldgetbeerorcoffeesoon')\n",
    "ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.1+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}