/ diplomatrix.py
diplomatrix.py
  1  import secrets
  2  import os
  3  
  4  # Define the characters that can be used in the one-time pad
  5  ALPHABET = 'üÜÁÍÓÚáíóúéèêëâàçîôûùïœÉÈÊËÂÀÇÎÔÛÙÏŒabcdefghijklmnñopqrstuvwxyzABCDEFGHIJKLMNÑOPQRSTUVWXYZ0123456789 .,()@"/=¿?!¡:;+-\''
  6  
  7  # Function to generate one-time pads
  8  def generate_one_time_pads(sheets, length):
  9      generated_sheets = []  # To store the filenames of generated sheets
 10      for sheet in range(sheets):
 11          filename = f"otp{sheet}.txt"
 12          generated_sheets.append(filename)
 13          with open(filename, "w") as f:
 14              for _ in range(length):
 15                  f.write(str(secrets.randbelow(len(ALPHABET))) + "\n")
 16      return generated_sheets
 17  
 18  # Function to load a one-time pad sheet from a file
 19  def load_sheet(filename):
 20      try:
 21          with open(filename, "r") as f:
 22              contents = f.read().splitlines()
 23          return contents
 24      except FileNotFoundError:
 25          print("Error: File not found.")
 26          return None
 27  
 28  # Function to get plaintext input from the user
 29  def get_plaintext():
 30      plaintext = input('Please type your message: ')
 31      return plaintext
 32  
 33  # Function to load content from a file
 34  def load_file(filename):
 35      try:
 36          with open(filename, "r") as f:
 37              contents = f.read()
 38          return contents
 39      except FileNotFoundError:
 40          print("Error: File not found.")
 41          return None
 42  
 43  # Function to save data to a file
 44  def save_file(filename, data):
 45      try:
 46          with open(filename, 'w') as f:
 47              f.write(data)
 48      except Exception as e:
 49          print(f"Error while saving: {e}")
 50  
 51  # Function to encrypt plaintext using a one-time pad sheet
 52  def encrypt(plaintext, sheet):
 53      ciphertext = ''
 54      for position, character in enumerate(plaintext):
 55          encrypted = (ALPHABET.index(character) + int(sheet[position])) % len(ALPHABET)
 56          ciphertext += ALPHABET[encrypted]
 57      return ciphertext
 58  
 59  # Function to decrypt ciphertext using a one-time pad sheet
 60  def decrypt(ciphertext, sheet):
 61      plaintext = ''
 62      for position, character in enumerate(ciphertext):
 63          decrypted = (ALPHABET.index(character) - int(sheet[position])) % len(ALPHABET)
 64          plaintext += ALPHABET[decrypted]
 65      return plaintext
 66  
 67  # Function to display the menu and handle user choices
 68  def menu():
 69      choices = ['1', '2', '3', '4']
 70      while True:
 71          print('What would you like to do?')
 72          print('1. Generate one-time pads')
 73          print('2. Encrypt a message')
 74          print('3. Decrypt a message')
 75          print('4. Quit the program')
 76          choice = input('Please type 1, 2, 3, or 4 and press Enter: ')
 77  
 78          if choice == '1':
 79              sheets = int(input('How many one-time pads would you like to generate? '))
 80              length = int(input('What will be your maximum message length? '))
 81              generated_sheets = generate_one_time_pads(sheets, length)
 82              print("Generated OTP sheets:")
 83              current_directory = os.getcwd()
 84              for sheet in generated_sheets:
 85                  print(os.path.join(current_directory, sheet))
 86  
 87          elif choice == '2':
 88              filename = input('Type in the filename of the OTP you want to use: ')
 89              sheet = load_sheet(filename)
 90              if sheet is not None:
 91                  plaintext = get_plaintext()
 92                  ciphertext = encrypt(plaintext, sheet)
 93                  output_filename = input('What will be the name of the encrypted file? ')
 94                  save_file(output_filename, ciphertext)
 95                  
 96                  # Display the location of the saved encrypted file
 97                  current_directory = os.getcwd()
 98                  encrypted_file_location = os.path.join(current_directory, output_filename)
 99                  print(f"The encrypted file has been saved at: {encrypted_file_location}")
100                  
101                  # Delete the used OTP sheet after encryption
102                  try:
103                      os.remove(filename)
104                      print(f"The OTP sheet '{filename}' has been deleted.")
105                  except Exception as e:
106                      print(f"Error while deleting OTP sheet: {e}")
107  
108          elif choice == '3':
109              filename = input('Type in the filename of the OTP you want to use: ')
110              sheet = load_sheet(filename)
111              if sheet is not None:
112                  ciphertext_filename = input('Type in the name of the file to be decrypted: ')
113                  ciphertext = load_file(ciphertext_filename)
114                  if ciphertext is not None:
115                      plaintext = decrypt(ciphertext, sheet)
116                      print('The message reads:')
117                      print('')
118                      print(plaintext)
119                      
120                      # Delete the used OTP sheet after decryption
121                      try:
122                          os.remove(filename)
123                          print(f"The OTP sheet '{filename}' has been deleted.")
124                      except Exception as e:
125                          print(f"Error while deleting OTP sheet: {e}")
126  
127          elif choice == '4':
128              exit()
129  
130          else:
131              print('Invalid choice. Please select a valid option.')
132  
133  # Start the menu loop
134  menu()
135