diff --git a/src/vigenere/break_vig.py b/src/vigenere/break_vig.py
index 2575592..5e532b7 100644
--- a/src/vigenere/break_vig.py
+++ b/src/vigenere/break_vig.py
@@ -1,25 +1,34 @@
-def probe(key):
+import argparse
+import re
+import vig
+import os
+
+
+def probe(key, substrings):
     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
+                count += 1
         key_hits[key+chr(x)] = count
-    best_val = [val for key, val in sorted(key_hits.items(), key = lambda ele: ele[1])][-1]
+    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}
 
 
-def crack(max_depth):
+def crack(substrings):
     tree = {"": 0}
     while tree:
-        curr = [key for key, val in sorted(tree.items(), key = lambda ele: ele[1])][-1]
-        if len(curr) == max_depth:
+        curr = [key for key, val in sorted(
+            tree.items(), key=lambda ele: ele[1])][-1]
+        print(curr)
+        if len(curr) == len(substrings[0]):
             return curr
         del tree[curr]
-        tree.update(probe(curr))
+        tree.update(probe(curr, substrings))
+
 
-import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument('FILE')
 parser.add_argument('--keylen', type=int, metavar="INT", required=True)
@@ -28,10 +37,5 @@ 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)]
 
-import os
 words = open(os.path.dirname(__file__)+"/words.txt", "r").read()
-
-import vig
-import re
-
-print(crack(args.keylen))
+print(crack(substrings))