#!/usr/bin/env python -i

import random
from fractions import gcd


class Coding:

	def encode(self, chars):
		nums = []
		for char in chars:
			char = char.upper()
			nums.append(ord(char) - ord('A'))
		return nums
		
	def decode(self, nums):
		chars = []
		for num in nums:
			chars.append(chr(num + ord('A')))
		return chars


class AffineCipher:
	m = 0
	k = [0, 0]
	ce = Coding() # The encoding engine

	def __init__(self, m = 26):
		self.keygen(m)
		
	def egcd(self,a, b):
	    if a == 0:
	        return (b, 0, 1)
	    else:
	        g, y, x = self.egcd(b % a, a)
	        return (g, x - (b // a) * y, y)

	def modinv(self, a, m):
	    g, x, y = self.egcd(a, m)
	    if g != 1:
	        raise Exception('modular inverse for '+a+' modulo '+m+
	        	' does not exist')
	    else:
	        return x % m
	
	def keygen(self, m = 26):
		self.m = m
		
		# This is a cryptographic safe pseudorandom number generator 
		# (PRG), in contrast to the build-in PRG.
		rand = random.SystemRandom()
		
		# Selects random numbers in the range [0, m-1], as long as the 
		# first one is coprime with m.
		a_ok = False
		while not a_ok:
			self.k[0] = rand.randrange(0, self.m)
			a_ok = (gcd(self.m, self.k[0]) == 1)
		self.k[1] = rand.randrange(0, self.m)
		
		return True
	
	def get_naive_security(self):
		# Approximately how many bits must be guessed in order to guess
		# the key? Note that this approximation is too high. Why?
		return 2 * self.m.bit_length()
		
	def encrypt(self, msg):
		# msg should be a string without any spaces. Case insensitive.
		chars = list(msg)
		nums = self.ce.encode(chars)
		for i in range(0, len(nums)):
			nums[i] = (self.k[0] * nums[i] + self.k[1]) % self.m
		chars = self.ce.decode(nums)
		return ''.join(chars)	
		
	def decrypt(self, c):
		# c should be a string without any spaces. Case insensitive.
		chars = list(c)
		nums = self.ce.encode(chars)
		for i in range(0, len(nums)):
			nums[i] = (self.modinv(self.k[0], self.m) * (nums[i] - 
				self.k[1])) % self.m
		chars = self.ce.decode(nums)
		return ''.join(chars)
		
	def bruteforce(self, c):
		# You need to complete this snippet yourself. Here's the
		# structure you should follow.
		
		# We need to generate all valid keys. Since we want to try all
		# keys, we should do it systematically, starting with the first
		# and ending with the last. The key generation routine is 
		# therefore not good -- it's a random algorithm, but we can 
		# copy some of the code from it.
		
		triedall = False
		while not triedall:
			# Generate the next key.
			# <insert your code here, updating k>
			
			# To avoid manual labour, search for the given string 
			# instead of printing everything.
			print self.k, self.decrypt(c)
			# Remember to set triedall = True when you have generated
			# the last key, which ought to be (25, 25).
			
		# For elegance: Reset k to the key it was before we started the
		# search.
		
print """
Usage: python -i affine.py
   or: ./affine.py (if you have set the executable flag)
	
The file provides the class AffineCipher, which you can use by typing

	cs = AffineCipher(<plaintext space size>)
	
If no argument is given, the plain text space is 26 by default. cs is
now an variable holding an instance of the cryptosystem. It contains
the following methods:

	keygen(m): Generates a new key with plaintext space size m
	encrypt(msg): Encrypts the string msg
	decrypt(c): Decrypts the ciphertext c
	get_naive_security(): Gives a crude approximation of the keylength
	                      in bits. The figure is too large, why?
	bruteforce(c): Does a bruteforce search for the key.
	               You need to implement this one yourself!
	               
Example

	cs = AffineCipher(26)
	c1 = cs.encrypt('HEI')
	c1
	cs.decrypt(c1)
	
Exercise: Decrypt the following ciphertext. You can assume that the
system is instantiated with m = 26, but you need to implement the
bruteforce method yourself.

	c1 = 'BILEEHPNZA'
	
It is known that the plaintext contains the substring 'GU'.
"""