break vig working

src/vigenere/break_vig.py

@@ -1,51 +1,24 @@
-#def probe(key):
-#    key_hits = {}
-#    for x in range(ord('a'), ord('z')+1):
-#        count = 0
-#        for substr in substrings:
-#            if re.search(vig.decrypt(key+chr(x), substr[0:len(key)+1]), words):
-#                count+=1
-#        key_hits[key+chr(x)] = count
-#    return ([val for key, val in sorted(key_hits.items(), key = lambda ele: ele[1])][-1], key_hits) 
-    
-
-def crack(current_key, max_depth):
+def probe(key):
     key_hits = {}
     for x in range(ord('a'), ord('z')+1):
         count = 0
         for substr in substrings:
-            if re.search(vig.decrypt(current_key+chr(x), substr[0:len(current_key)+1]), words):
+            if re.search(vig.decrypt(key+chr(x), substr[0:len(key)+1]), words):
                 count+=1
-        key_hits[current_key+chr(x)] = count
-    best = [key for key, val in sorted(key_hits.items(), key = lambda ele: ele[1])][-1]
-    print("best = {} ".format(best))
-
-    if len(current_key)+1 < max_depth:
-        crack(best, max_depth)
-
-    #bestkeys = []
-    #for key, val in key_hits.items():
-    #    if val == best:
-    #        bestkeys.append(key)
-    #print(bestkeys)
+        key_hits[key+chr(x)] = count
+    best_val = [val for key, val in sorted(key_hits.items(), key = lambda ele: ele[1])][-1]
+    return {key: val for key, val in key_hits.items() if val == best_val}
 
-    #bestprobes = {bestkeys[0]: 0}
-    #if len(bestkeys) > 1:
-    #    for key in bestkeys:
-    #        print("probing {} ... ".format(key), end='')
-    #        num = probe(key)[0]
-    #        print(num)
-    #        if(num > list(bestprobes.values())[0]):
-    #            bestprobes = {}
-    #            bestprobes[key] = num
-    #        if(num == list(bestprobes.values())[0]):
-    #            bestprobes[key] = num
-    #    print("best probes: {}".format(bestprobes))
-
-    #for item in bestprobes.items():
-    #    crack(item[0], max_depth)
-    #return crack(bestprobe[0], max_depth)
 
+def crack(max_depth):
+    tree = {"": 0}
+    while tree:
+        curr = [key for key, val in sorted(tree.items(), key = lambda ele: ele[1])][-1]
+        print(curr)
+        if len(curr) == max_depth:
+            return curr
+        del tree[curr]
+        tree.update(probe(curr))
 
 import argparse
 parser = argparse.ArgumentParser()
@@ -59,9 +32,7 @@ substrings = [(t[i:i+args.keylen]) for i in range(0, len(t), args.keylen)]
 import os
 words = open(os.path.dirname(__file__)+"/words.txt", "r").read()
 
-
-
 import vig
 import re
 
-crack("", args.keylen)
+crack(args.keylen)