diff --git a/src/vigenere/break_vig.py b/src/vigenere/break_vig.py index d380b9a..a6e8674 100644 --- a/src/vigenere/break_vig.py +++ b/src/vigenere/break_vig.py @@ -28,13 +28,14 @@ def crack(substrings, words): tree.update(probe(curr, substrings, words)) -parser = argparse.ArgumentParser() -parser.add_argument('FILE') -parser.add_argument('--keylen', type=int, metavar="INT", required=True) -args = parser.parse_args() +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('FILE') + parser.add_argument('--keylen', type=int, metavar="INT", required=True) + args = parser.parse_args() -t = ''.join([x for x in open(args.FILE, "r").read().lower() if x.isalpha()]) -substrings = [(t[i:i+args.keylen]) for i in range(0, len(t), args.keylen)] -words = open(os.path.dirname(__file__)+"/words.txt", "r").read() + t = ''.join([x for x in open(args.FILE, "r").read().lower() if x.isalpha()]) + substrings = [(t[i:i+args.keylen]) for i in range(0, len(t), args.keylen)] + words = open(os.path.abspath(os.path.dirname(__file__))+"/words.txt", "r").read() -print(crack(substrings, words)) + print(crack(substrings, words))