#!/usr/bin/env python

# library of functions to demonstrate simple ciphers and 
# their weaknesses

import string
import random
import time
random.seed(time.time())

# alphabet defined globally for all functions
alphabet="abcdefghijklmnopqrstuvwxyz"

def caesar(input,key,decrypt=False):
  # Implements a simple Ceasar cipher, or ROT N for values
  # of N from 1 to 25. Can encrypt or decrypt.
  # Performs all manipulations on lower case letters.
  input=string.lower(input)
  output="" # empty string
  # ensure key is a number between zero and length of alphabet - 1.
  size_alph=len(alphabet) # 26 for lowercase Roman alphabet
  key=key%size_alph

  # process all characters in turn
  for char in input:
    letternum=alphabet.rfind(char) # gets 1 for a, 2 for b etc. -1 for punctuation
    if letternum >= 0 and letternum < size_alph: 
      if decrypt:
        opos=(letternum + size_alph - key) % size_alph
      else: # encrypt
        opos=(letternum + size_alph + key) % size_alph
      # add letter at end of ouput string
      output=output+alphabet[opos]
    else: # pass punctuation straight through
      output=output+char
  return output

def gen_subst_key():
  # subst gives a set of 26 rotations for letters a-z so that
  # each letter substitutes to a random other, and every substitute
  # letter is used for exactly one original letter. 
  # This is achieved by shuffling randomly a cloned
  # copy of the alphabet as one could a pack of cards. 
  # 
  # In this 2nd version a letter can substitute for itself.
  subst=alphabet[:] # full slice to clone it. We don't want a reference.
  # shuffle the alphabet 3 times to generate substitution key
  subst=shuffle(subst) 
  return subst 

def shuffle(str):
  ''' shuffles string letters swapping each letter 
  for one in a random position anywhere in the string. The
  same set of letters is output as input, but in a different order.
  The shuffle is repeated 3 times. '''
  l=list(str) # lists are mutable, strings are not
  lenstr=len(str)
  for j in range(3):
    for i in range(lenstr): 
      randpos=random.randint(0,lenstr-1)
      l[i],l[randpos]=l[randpos],l[i] # don't need temp var for a Python swap
  # convert list of chars back into string
  shuffled=''
  for letter in l:
    shuffled+=letter
  return shuffled

def gen_transp_key():
  # define positions in block to be used
  blockpos="0123456789"
  transpkey=shuffle(blockpos)
  return transpkey 

def pad(input,blocksize,messagesize,pad=True):
  # pads input to multiple of blocksize using
  # random chars taken from alphabet
  # set pad parameter to false to strip padding
  if not pad: # strip padding
    return input[0:messagesize] # Python slice operator
  else: # pad message
    # first clone input
    output=input[:] # slice of entire string clones it 
    while len(output) % blocksize != 0:
      # add random chars to output until multiple of blocksize
      randpos=random.randint(0,len(alphabet)-1)
      output=output+alphabet[randpos]
    return output

def subst(input,key,decrypt=False):
  # performs a substitution cypher rotating
  # each letter in input to the value at
  # the position of the input letter in
  # the keystring, such that a substitutes to the first
  # letter in key, b substitutes to the second etc.
  
  # Performs all manipulations on lower case letters.
  input=string.lower(input)
  output="" # starts empty, adds a char at a time
  for char in input: # loop once for each input char
    if char in alphabet:  
      if decrypt:
        # find position of char in key
        keypos=key.rfind(char)
        output=output+alphabet[keypos]
      else: # must be encrypting
        alphpos=alphabet.rfind(char) # 0 for a, 1 for b etc.
        output=output+key[alphpos]
    else:
      output=output+char
  return output 

def trans(input,key,decrypt=False):
  # performs transposition cipher encrypt or decrypt on input
  # input must be padded to multiple of block and keysize
  blocksize=len(key)
  if len(input) % len(key) != 0:
    raise ValueError # input must be multiple of blocksize
  output="" # start as empty string 
  sortkey="0123456789"

  # mark start positions of each transpose block in blockstarts list
  blockstarts=[]
  for i in range(0,len(input)):
    if i%blocksize==0:
      blockstarts.append(i) # [0,10,20,30 ... ]
  
  # process all blocks
  for bstart in blockstarts:
    inblock=input[bstart:bstart+blocksize] # slice input block
    if decrypt: 
      for num in sortkey:
        posnext=key.rfind(num)
        output=output+inblock[posnext]
    else: # decrypt
      for num in key:
        posnext=sortkey.rfind(num)
        output=output+inblock[posnext]
  return output

def block(input,key,decrypt=False):
  # performs simple block cipher by chaining substitution and tranposition ciphers
  skey,tkey=key.split(",") # extract substitution and transposition keys
  blocksize=len(tkey)
  if decrypt:
    # split message length from ciphertext
    slen,ciphertext=input.split(",") # extract substitution and transposition keys
    ilen=int(slen)
    # untranspose 
    midpad=trans(ciphertext,tkey,decrypt=True)
    # remove padding
    midtext=pad(midpad,10,ilen,pad=False)
    # unsusbstitute
    plaintext=subst(midtext,skey,decrypt=True)
    return plaintext 
  else: # encrypt
    # get message length
    ilen=len(input)
    slen=str(ilen)
    # substitute
    midtext=subst(input,skey)
    # transpose
    padmid=pad(midtext,10,len(midtext))
    ciphertext=trans(padmid,tkey)
    # prepend message length string to ciphertext
    output=slen+","+ciphertext
    return output

def main():
  # run all defined test cases
  allcases=["caesar", "gen_subst_key", "gen_transp_key", "pad",
    "subst", "trans", "block" ]
  for case in allcases:
    print "test case: "+case
    test(case)
    raw_input("press enter to continue")
    print 

def test(case="caesar"):

  if case == "caesar": 
    # ceasar test case
    key=17
    plaintext="this is a message"
    ciphertext=caesar(plaintext,key)
    decrypt=caesar(ciphertext,key,decrypt=True)
    print "plaintext: "+plaintext
    print "ciphertext: "+ciphertext
    print "decrypt: "+decrypt 

  elif case == "gen_subst_key":
    print alphabet
    print gen_subst_key()
  
  elif case == "gen_transp_key":
    blockpos="0123456789"
    print blockpos 
    print gen_transp_key()

  elif case == "pad":
    message="this is a message not of multiple 10"
    padmessage=pad(message,10,len(message))
    stripmessage=pad(padmessage,10,len(message),pad=False)
    print "message: "+message 
    print "length of message: "+str(len(message))
    print "padmessage: "+padmessage 
    print "stripmessage: "+stripmessage 

  elif case == "trans":
    plaintext="prime numbers may contain the secrets of the universe"
    key=gen_transp_key()
    padplain=pad(plaintext,10,len(plaintext))
    ciphertext=trans(padplain,key)
    decryptpad=trans(ciphertext,key,decrypt=True)
    decrypt=pad(decryptpad,10,len(plaintext),pad=False)
    print "key: "+key
    print "plaintext: "+plaintext
    print "length of plaintext: "+str(len(plaintext))
    print "padplain: "+padplain
    print "ciphertext: "+ciphertext
    print "decryptpad: "+decryptpad 
    print "decrypt: "+decrypt 

  elif case == "subst":
    key=gen_subst_key()
    plaintext="the quick brown fox jumps over the lazy dog"
    ciphertext=subst(plaintext,key)
    decrypt=subst(ciphertext,key,decrypt=True)
    print "key: "+key
    print "plaintext: "+plaintext
    print "ciphertext: "+ciphertext
    print "decrypt: "+decrypt 

  elif case == "block":
    # generate combined key
    skey=gen_subst_key()
    tkey=gen_transp_key()
    key=skey+","+tkey
    plaintext="the quick brown fox jumps over the lazy dog"
    ciphertext=block(plaintext,key)
    decrypt=block(ciphertext,key,decrypt=True)
    print "key: "+key
    print "plaintext: "+plaintext
    print "ciphertext: "+ciphertext
    print "decrypt: "+decrypt 

# run all test cases if module is run as a program  
if __name__ == "__main__":
  main()