import math

def main():
    message = input('Enter message: ')
    key = int(input('Enter key [2-%s]: ' % (len(message) - 1)))
    mode = input('Encryption/Decryption [e/d]: ')

    if mode.lower().startswith('e'):
        text = encryptMessage(key, message)
    elif mode.lower().startswith('d'):
        text = decryptMessage(key, message)

    # Append pipe symbol (vertical bar) to identify spaces at the end.
    print('Output:\n%s' %(text + '|'))

def encryptMessage(key, message):
    """
    >>> encryptMessage(6, 'Harshil Darji')
    'Hlia rDsahrij'
    """
    cipherText = [''] * key
    for col in range(key):
        pointer = col
        while pointer < len(message):
            cipherText[col] += message[pointer]
            pointer += key
    return ''.join(cipherText)

def decryptMessage(key, message):
    """
    >>> decryptMessage(6, 'Hlia rDsahrij')
    'Harshil Darji'
    """
    numCols = math.ceil(len(message) / key)
    numRows = key
    numShadedBoxes = (numCols * numRows) - len(message)
    plainText = [""] * numCols
    col = 0; row = 0;

    for symbol in message:
        plainText[col] += symbol
        col += 1

        if (col == numCols) or (col == numCols - 1) and (row >= numRows - numShadedBoxes):
            col = 0
            row += 1

    return "".join(plainText)

if __name__ == '__main__':
    import doctest
    doctest.testmod()
    main()