######################################################
#                                                    #
#  rollingcypher.py - a rolling substitution cypher  #
#                                                    #
#  A more complex version of the ceaser cypher.      #
#                                                    #
#  Every string is divided up into arbitrary length  #
#  blocks, with four padded characters on the end.   #
#  This padding is (a) the shift to use, and         #
#  (b) the length of the next segment. This extra    #
#  info is shuffled with the text before it, and     #
#  provides the key to the next segment.             #
#                                                    #
#  The 4-character key (unencrypted) is required to  #
#  decrypt the first segment, and thus the rest.     #
#                                                    #
######################################################

from random import *

#Original: [" ", "A", "a", "B", "b", "C", "c", "D", "d", "E", "e", "F", "f", "G", "g", "H", "h", "I", "i", "J", "j", "K", "k", "L", "l", "M", "m", "N", "n", "O", "o", "P", "p", "Q", "q", "R", "r", "S", "s", "T", "t", "U", "u", "V", "v", "W", "w", "X", "x", "Y", "y", "Z", "z", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "~", "`", "!", "@", "#", "$", "%", "^", "&", "*", "(", ")", "-", "_", "=", "+", "[", "]", "{", "}", "\\", "|", ";", ":", "'", "\"", ",", "<", ".", ">", "/", "?"]
characterlist = ['j', '+', '?', '$', 'a', 'c', 'g', 'T', '\\', 's', '"', ' ', '@', '0', 't', 'p', 'L', ',', 'u', '/', 'Q', '`', 'm', 'n', 'F', '}', 'U', '#', 'l', ':', '[', '>', '6', 'w', 'o', '.', '(', '8', '9', '*', 'C', '=', 'K', '%', 'r', '!', 'v', '7', 'k', '2', '3', 'H', 'P', 'x', 'z', 'M', 'Y', 'O', 'Z', "'", 'B', '&', '~', 'V', 'h', '<', 'W', 'S', '^', 'b', 'X', '_', 'R', ')', 'i', 'D', 'E', 'y', ']', '4', 'G', '|', 'A', ';', 'N', 'e', 'I', '1', 'J', '5', 'q', 'f', '-', '{', 'd', '0']

def encrypt(cleartext, key = "0000"):
    """
    Encrypts the string. Key defaults to 0000, and if it is 0000, will generate a random key.

    Cleartext can be any string. Key must be 4 characters, padded with 0s if nessecary, shift then segment length
    The function returns (initial key, encrypted text)
    """

    ciphertext = ""
    shift  = 0
    seglen = 0
    if key == "0000":
        shift = randint(1,len(characterlist)-1)
        seglen = randint(1,30)

        if shift < 10:  shiftstr = "0"+str(shift)
        else:           shiftstr = str(shift)
        if seglen < 10: seglenstr = "0"+str(seglen)
        else:           seglenstr = str(seglen)
        key = shiftstr + seglenstr
    else:
        shift = int(key[0:2])
        seglen = int(key[2:4])

    while len(cleartext) > 0:
        #loop through each segment, split it off the cleartext, encrypt, then attach it to the ciphertext with key info
        segment = cleartext[:seglen]
        cleartext = cleartext[seglen:]

        #Calculate the next shift so it can be attached onto the end of the cleartext strip
        newshift  = randint(1,len(characterlist)-1)
        newseglen = randint(1,30)
        if newshift < 10: shiftstr = "0"+str(newshift)
        else:             shiftstr = str(newshift)
        if newseglen < 10: seglenstr = "0"+str(newseglen)
        else:              seglenstr = str(newseglen)

        segment = segment + shiftstr + seglenstr

        #replace each character by shifting x items through the character list
        for character in segment:
            pos = 0
            for position, item in enumerate(characterlist):
                if item == character:
                    pos = position
                    break

            newpos = pos + shift
            if newpos >= len(characterlist):
                newpos -= len(characterlist)

            ciphertext += characterlist[newpos]

        shift = newshift
        seglen = newseglen

    return (ciphertext, key)
#End encrypt

def decrypt(ciphertext, key):
    """
    Decrypts by reversing the above method. An initial key is nessecary

    Ciphertext is any string encrypted with the above function. Key must be the four character key used to encrypt
    The function returns the cleartext version
    """

    cleartext = ""
    shift  = int(key[0:2])
    seglen = int(key[2:4])

    while len(ciphertext) > 0:
        #loop through each segment, split off the ciphertext, decrypt, split the key info, then attach to cleartext
        segment    = ciphertext[:seglen+4]
        ciphertext = ciphertext[seglen+4:]

        for character in segment:
            pos = 0
            for position, item in enumerate(characterlist):
                if item == character:
                    pos = position
                    break

            newpos = pos - shift
            if newpos < 0:
                newpos += len(characterlist)

            cleartext += characterlist[newpos]

        #split off key info for the next iteration
        shift  = int(cleartext[-4:-2])
        seglen = int(cleartext[-2:])
        cleartext = cleartext[:-4]

    return cleartext
#End decrypt