Fix major bugs, major improvements in flag tasks and improve error handling

Author: disalad

project/__main__.py

@@ -6,11 +6,11 @@
     description='A tool to Encrypt, Decrypt texts with a substitution cipher given as the key')
 
 # -e & --encrypt
-parser.add_argument('-e', '--encrypt', nargs='+',
-                    metavar='PLAIN', help='Encrypt the plain text', type=str)
+parser.add_argument('-e', '--encrypt', nargs='?',
+                    metavar='PLAIN', help='Encrypt the plain text', type=str, default=False, const=True)
 # -d & --decrypt
-parser.add_argument('-d', '--decrypt', nargs='+', metavar='CIPHER',
-                    help='Decrypt the cipher text', type=str)
+parser.add_argument('-d', '--decrypt', nargs='?', metavar='CIPHER',
+                    help='Decrypt the cipher text', type=str, default=False, const=True)
 # -k & --key
 parser.add_argument('-k', '--key', nargs=1, metavar='KEY',
                     help='Represents the number of places for the shift', type=int, default=False)
@@ -20,34 +20,31 @@
 parser.add_argument('-o', '--output', nargs=1, type=str,
                     help='Output file directory')
 # -s & --source
-parser.add_argument('-sf', '--sourceFile', action='store_true')
+parser.add_argument('-sf', '--sourceFile', nargs=1, metavar='PATH',
+                    help='Path to the file containing message to cipher', type=str)
 
 # -b & --brute
-parser.add_argument('-br', '--bruteforce', nargs='+',
-                    metavar='PLAIN', help='Cipher text to crack', type=str)
+parser.add_argument('-br', '--bruteforce', nargs='?',
+                    metavar='PLAIN', help='Cipher text to crack', type=str, default=False, const=True)
 
 # -fr & --freq-analysis
-parser.add_argument('-fr', '--freq-analysis', nargs='+', metavar='CIPHER',
-                    help='Plain text to crack by analysing', type=str)
+parser.add_argument('-fr', '--freq-analysis', nargs='?', metavar='CIPHER',
+                    help='Plain text to crack by analysing', type=str, default=False, const=True)
 
 # -w & --wordlist
 parser.add_argument('-w', '--wordlist', nargs=1, metavar='WORDLIST',
                     help='Path to the wordlist containing language words', type=str)
 
 # -bfr & --brute-frequency
-parser.add_argument('-bfr', '--brute-frequency', nargs='+', metavar='CIPHER',
-                    help='Cipher to crack and analyse output for the text containing words in a wordlist')
+parser.add_argument('-bfr', '--brute-frequency', nargs='?', metavar='CIPHER',
+                    help='Cipher to crack and analyse output for the text containing words in a wordlist', default=False, const=True)
 
 args = parser.parse_args()
 
 
 def handle_args():
     try:
-        # If key isn't in the args,
-        if not args.key and (args.encrypt or args.decrypt):
-            raise NameError("Key is not defined")
-        # If brute frequency flag used,
-        elif args.brute_frequency:
+        if args.brute_frequency:
             brute_freq_handler(args)
         # If frequency flag used,
         elif args.freq_analysis:

project/crack.py

@@ -1,14 +1,15 @@
 from .crypt import ALPHEBET, decrypt
 from .utils import analyse_text
 from operator import itemgetter
+from colorama import Fore, Style
 import re
 
 
 def bruteforce(text):
     cracked = []
     # Iterate from 1 to length of the ALPHEBET
     for i in range(1, len(ALPHEBET)):
-        cracked.append('Result with key ' + str(i) + ' ' + decrypt(text, i))
+        cracked.append(f'{Fore.GREEN}{str(i)}{Style.RESET_ALL} : {decrypt(text, i)}')
     return cracked
 
 
@@ -25,5 +26,5 @@ def brute_freq(text, wordlist):
         avg = float(matches) / len(text.split(' ')) * 100
         analysed[text] = avg
     # Sort the analysed texts by average
-    analysed = dict(sorted(analysed.items(), key=itemgetter(1), reverse=True))
-    return analysed
+    sorted_analysed = dict(sorted(analysed.items(), key=itemgetter(1), reverse=True))
+    return [sorted_analysed, analysed]

project/crypt.py

@@ -21,10 +21,7 @@ def encrypt(plain_text, key):
                 ciphered = ALPHEBET[ALPHEBET.find(lower_i) + key]
             # If the letter is an uppercase,
             # cipher letter will be uppercased too
-            if i.isupper():
-                cipher += ciphered.upper()
-            else:
-                cipher += ciphered
+            cipher += ciphered.upper() if i.isupper() else ciphered
         # If letter doesn't occur in the ALPHEBET,
         # the plain letter will be added to the cipher text
         else:
@@ -47,10 +44,7 @@ def decrypt(cipher_text, key):
                 decrypted = ALPHEBET[ALPHEBET.find(lower_i) - key]
             # If the letter is an uppercase,
             # decrypted letter will be uppercased too
-            if i.isupper():
-                plain += decrypted.upper()
-            else:
-                plain += decrypted
+            plain += decrypted.upper() if i.isupper() else decrypted
         # If letter doesn't occur in the ALPHEBET,
         # the plain letter will be added to the plain text
         else:

project/handlers.py

@@ -1,18 +1,13 @@
 from .crack import bruteforce, brute_freq
-from .utils import write_file, analyse_text
+from .utils import write_file, analyse_text, get_text
 from .crypt import encrypt, decrypt
 from colorama import Fore, Style
 import re
 
 
 def bruteforce_handler(args):
-    c_value = ' '.join(args.bruteforce)
-    cracked = []
-    if args.sourceFile:
-        f = open(c_value, 'r').read()
-        cracked = bruteforce(f)
-    else:
-        cracked = bruteforce(c_value)
+    # Get all possibilities of the cipher text
+    cracked = bruteforce(get_text(args, 'bruteforce'))
     # Write to the file,
     # If an output path specified
     if args.output:
@@ -24,18 +19,9 @@ def bruteforce_handler(args):
 
 def encrypt_handler(args):
     # Extract key as an int from args
-    key = int(''.join(map(str, args.key)))
-    key *= -1 if args.backwards else 1
-    # Get plain text
-    e_value = ' '.join(args.encrypt)
-    # If source presents,
-    # file content will be encrypted
-    cipher = ''
-    if args.sourceFile:
-        f = open(e_value, 'r').read()
-        cipher = encrypt(f, key)
-    else:
-        cipher = encrypt(e_value, key)
+    key = int(''.join(map(str, args.key))) * -1 if args.backwards else 1
+    # Get Cipher text
+    cipher = encrypt(get_text(args, 'encrypt'), key)
     # Write to the file,
     # If an output path specified
     if args.output:
@@ -46,18 +32,9 @@ def encrypt_handler(args):
 
 def decrypt_handler(args):
     # Extract key as an int from args
-    key = int(''.join(map(str, args.key)))
-    key *= -1 if args.backwards else 1
-    # Get Cipher text
-    d_value = ' '.join(args.decrypt)
-    # If source presents,
-    # file content will be decrypted
-    plain = ''
-    if args.sourceFile:
-        f = open(d_value, 'r').read()
-        plain = decrypt(f, key)
-    else:
-        plain = decrypt(d_value, key)
+    key = int(''.join(map(str, args.key))) * -1 if args.backwards else 1
+    # Get plain text
+    plain = decrypt(get_text(args, 'decrypt'), key)
     # Write to the file,
     # If an output path specified
     if args.output:
@@ -70,10 +47,8 @@ def freq_analysis_handler(args):
     # Check for missing arguments
     if not args.wordlist:
         raise ValueError("No Wordlist Found")
-    text = ''
-    a_value = ' '.join(args.freq_analysis)
     # Filter letters, numbers and spaces
-    text = re.sub(r'[^A-Za-z0-9 ]+', '', a_value)
+    text = re.sub(r'[^A-Za-z0-9 ]+', '', get_text(args, 'freq_analysis'))
     # Analyse the frequency
     matches = analyse_text(text, ''.join(args.wordlist))
     avg = float(matches) / len(text.split(' ')) * 100
@@ -84,11 +59,14 @@ def freq_analysis_handler(args):
 
 
 def brute_freq_handler(args):
-    c_value = ' '.join(args.brute_frequency)
+    # Check for missing arguments
+    if not args.wordlist:
+        raise ValueError("No Wordlist Found")
     # Get cipher text
-    cipher = open(c_value, 'r').read() if args.sourceFile else c_value
+    cipher = get_text(args, 'brute_frequency')
     # Get analysed disctionary
-    analyzed = brute_freq(cipher, ''.join(args.wordlist))
+    bruteforced = brute_freq(cipher, ''.join(args.wordlist))
+    analyzed = bruteforced[0]
     # Write to the file,
     # If an output path specified
     if args.output:
@@ -105,7 +83,8 @@ def brute_freq_handler(args):
         found = False
         for i in analyzed:
             if analyzed[i] > 60:
-                print(f'{Fore.GREEN}{round(analyzed[i], 1)}%{Style.RESET_ALL} : {i}')
+                print(
+                    f'{Fore.GREEN}{round(analyzed[i], 1)}%{Style.RESET_ALL}: {Fore.CYAN}{list(bruteforced[1].keys()).index(i)}{Style.RESET_ALL} : {i}')
                 found = True
         if not found:
             print(

project/utils.py

@@ -29,3 +29,10 @@ def analyse_text(text, wordlist_path):
         if word in wordlist:
             matches += 1
     return matches
+
+
+def get_text(args, prop):
+    if args.sourceFile:
+        return open(' '.join(args.sourceFile), 'r').read()
+    else:
+        return ''.join(vars(args)[prop])