# Skeleton Program code for the AQA COMP1 Summer 2013 examination
# this code should be used in conjunction with the Preliminary Material
# written by the AQA COMP1 Programmer Team
# developed in the Python 3.0.1 programming environment

def DisplayMenu():
  print()
  print('PLAINTEXT OPTIONS')
  print('  a.  Get plaintext from user')
  print('  b.  Display plaintext')
  print('CIPHERTEXT OPTIONS')
  print('  d.  Get ciphertext from user')
  print('  e.  Display ciphertext')
  print('ENCRYPT')
  print('  g.  Caesar cipher')
  print('  h.  Rail fence')
  print('DECRYPT')
  print('  j.  Caesar cipher')
  print('  k.  Rail fence')
  print('STEGANOGRAPHY')
  print('  n.  nth character (text from file)')
  print()
  print('Please select an option from the above list (or enter q to quit): ', end='')

def GetMenuChoice():
  MenuChoice = input()
  print()
  return MenuChoice

def GetTextFromUser():
  TextFromUser = input('Please enter the text to use: ')
  return TextFromUser

def GetPositionsToUse():
  StartPosition = int(input('Please enter the start position to use in the file: '))
  EndPosition = int(input('Please enter the end position to use in the file: '))
  return StartPosition, EndPosition

def GetTextFromFile(StartPosition, EndPosition):
  CurrentFile = open('diary.txt', 'rb')
  for Count in range(1, StartPosition):
    ByteFromFile = CurrentFile.read(1) 
  TextFromFile = ''
  for Count in range(StartPosition, EndPosition + 1):
    ByteFromFile = CurrentFile.read(1)
    TextFromFile += ByteFromFile.decode('ASCII')
  CurrentFile.close()
  return TextFromFile

def GetKeyForCaesarCipher():
  Key = int(input('Enter the amount that shifts the plaintext alphabet to the ciphertext alphabet: '))
  return Key

def GetTypeOfCharacter(ASCIICode):
  if ASCIICode >= ord('A') and ASCIICode <= ord('Z'):
    TypeOfCharacter = 'Upper'
  elif ASCIICode >= ord('a') and ASCIICode <= ord('z'):
    TypeOfCharacter = 'Lower'
  else:
    TypeOfCharacter = 'Other'
  return TypeOfCharacter

def ApplyShiftToASCIICodeForCharacter(ASCIICode, AmountToShift):
  TypeOfCharacter = GetTypeOfCharacter(ASCIICode)
  if TypeOfCharacter != 'Other':
    if TypeOfCharacter == 'Upper':
      NewASCIICode = ((26 + ASCIICode - ord('A') + AmountToShift) % 26) + ord('A')
    else:
      NewASCIICode = ((26 + ASCIICode - ord('a') + AmountToShift) % 26) + ord('a')
  else:
    NewASCIICode = ASCIICode
  return NewASCIICode

def UseCaesarCipher(OriginalText, AmountToShift):
  ChangedText = ''
  for Count in range(0, len(OriginalText)):
    ASCIICode = ord(OriginalText[Count])
    ASCIICode = ApplyShiftToASCIICodeForCharacter(ASCIICode, AmountToShift)
    ChangedText += chr(ASCIICode)
  return ChangedText

def GetSizeOfRailFence():
  SizeOfRailFence = int(input('Enter the number of lines in the rail fence: '))
  return SizeOfRailFence

def EncryptUsingRailFence(OriginalText, SizeOfRailFence):
  Ciphertext = ''
  for Count1 in range(0, SizeOfRailFence):
    Count2 = Count1
    while Count2 < len(OriginalText):
      Ciphertext += OriginalText[Count2]
      Count2 = Count2 + SizeOfRailFence
  return Ciphertext

def DecryptUsingRailFence(Ciphertext, SizeOfRailFence):
  Plaintext = ''
  NoOfCiphertextCharacters = len(Ciphertext)
  # The ciphertext was created from a visualisation of the plaintext as
  # a two-dimensional grid of characters with no of rows = size of rail fence
  NoOfRows = SizeOfRailFence
  NoOfColumns = NoOfCiphertextCharacters // SizeOfRailFence
  # If NoOfCiphertextCharacters divides exactly all rows will be full
  # otherwise the last column will be incomplete and NoOfColumns will not include last column
  if (NoOfCiphertextCharacters % SizeOfRailFence) != 0:
    NoOfColumns += 1
  # Calculate row no of last full row, 0 means every row full
  LastFullRowNo = NoOfCiphertextCharacters % SizeOfRailFence
  NoOfCiphertextCharactersProcessed = 0
  for i in range (0, NoOfColumns):  # Work along the columns building the plaintext a column at a time
    AmountToReduceNoOfColumnsTimesjBy = 0
    for j in range (0, NoOfRows):  # Work down the rows building the plaintext
      if LastFullRowNo != 0:  # Last column doesn't have a character in every row
        if j > LastFullRowNo: # There are shorter rows to skip
          AmountToReduceNoOfColumnsTimesjBy += 1
      # NoOfColumns * j locates in ciphertext beginning of each row
      BeginningOfNextRowIndex = NoOfColumns * j - AmountToReduceNoOfColumnsTimesjBy
      PositionOfNextCharacter = BeginningOfNextRowIndex + i
      NoOfCiphertextCharactersProcessed += 1
      if NoOfCiphertextCharactersProcessed <= NoOfCiphertextCharacters:
        Plaintext += Ciphertext[PositionOfNextCharacter]
  return Plaintext

def GetValueForN():
  N = int(input('Enter the value of n: '))
  return N

def EveryNthCharacterSteganography(StartPosition, EndPosition, N):
  CurrentPosition = StartPosition
  HiddenMessage = ''
  while CurrentPosition <= EndPosition:
    HiddenMessage += GetTextFromFile(CurrentPosition, CurrentPosition)
    CurrentPosition = CurrentPosition + N
  return HiddenMessage

def DisplayPlaintext(TextToDisplay):
  print()
  print('The plaintext is: ', end='')
  print(TextToDisplay)

def DisplayCiphertext(TextToDisplay):
  print()
  print('The ciphertext is: ', end='')
  print(TextToDisplay)

if __name__ == '__main__':
  Ciphertext = ''
  Plaintext = ''
  Choice = '0'
  StartPosition = 0
  EndPosition = 0
  while Choice != 'q':
    DisplayMenu()
    Choice = GetMenuChoice()
    if Choice == 'a':
      Plaintext = GetTextFromUser()
    elif Choice == 'b':
      DisplayPlaintext(Plaintext)
    elif Choice == 'd':
      Ciphertext = GetTextFromUser()
    elif Choice == 'e':
      DisplayCiphertext(Ciphertext)
    elif Choice == 'g':
      DisplayPlaintext(Plaintext)
      AmountToShift = GetKeyForCaesarCipher()
      Ciphertext = UseCaesarCipher(Plaintext, AmountToShift)
      DisplayCiphertext(Ciphertext)
    elif Choice == 'h':
      DisplayPlaintext(Plaintext)
      SizeOfRailFence = GetSizeOfRailFence()
      Ciphertext = EncryptUsingRailFence(Plaintext, SizeOfRailFence)
      DisplayCiphertext(Ciphertext)
    elif Choice == 'j':
      DisplayCiphertext(Ciphertext)
      AmountToShift = -GetKeyForCaesarCipher()
      Plaintext = UseCaesarCipher(Ciphertext, AmountToShift)
      DisplayPlaintext(Plaintext)
    elif Choice == 'k':
      DisplayCiphertext(Ciphertext)
      SizeOfRailFence = GetSizeOfRailFence()
      Plaintext = DecryptUsingRailFence(Ciphertext, SizeOfRailFence)
      DisplayPlaintext(Plaintext)
    elif Choice == 'n':
      StartPosition, EndPosition = GetPositionsToUse()
      N = GetValueForN()
      Plaintext = EveryNthCharacterSteganography(StartPosition, EndPosition, N)
      DisplayPlaintext(Plaintext)
    if Choice != 'q':
      input('Press enter key to continue')