Merge pull request #60 from dxa4481/dev

Dev

tests.py

@@ -0,0 +1,26 @@
+import unittest
+import os
+from truffleHog import truffleHog
+
+
+class TestStringMethods(unittest.TestCase):
+
+    def test_shannon(self):
+        random_stringB64 = "ZWVTjPQSdhwRgl204Hc51YCsritMIzn8B=/p9UyeX7xu6KkAGqfm3FJ+oObLDNEva"
+        random_stringHex = "b3A0a1FDfe86dcCE945B72" 
+        self.assertGreater(truffleHog.shannon_entropy(random_stringB64, truffleHog.BASE64_CHARS), 4.5)
+        self.assertGreater(truffleHog.shannon_entropy(random_stringHex, truffleHog.HEX_CHARS), 3)
+
+    def test_cloning(self):
+        project_path = truffleHog.clone_git_repo("https://github.com/dxa4481/truffleHog.git")
+        license_file = os.path.join(project_path, "LICENSE")
+        self.assertTrue(os.path.isfile(license_file))
+
+    def test_unicode_expection(self):
+        try:
+            truffleHog.find_strings("https://github.com/dxa4481/tst.git")
+        except UnicodeEncodeError:
+            self.fail("Unicode print error")
+
+if __name__ == '__main__':
+    unittest.main()

truffleHog/truffleHog.py

@@ -70,15 +70,40 @@ class bcolors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
-def find_strings(git_url, printJson=False):
+def clone_git_repo(git_url):
     project_path = tempfile.mkdtemp()
     Repo.clone_from(git_url, project_path)
+    return project_path
+
+def print_results(printJson, commit_time, branch_name, prev_commit, printableDiff):
+    if printJson:
+        print(json.dumps(output, sort_keys=True, indent=4))
+    else:
+        if sys.version_info >= (3, 0):
+            dateStr = "{}Date: {}{}".format(bcolors.OKGREEN, commit_time, bcolors.ENDC)
+            print(dateStr)
+            branchStr = "{}Branch: {}{}".format(bcolors.OKGREEN, branch_name, bcolors.ENDC)
+            print(branchStr)
+            commitStr = "{}Commit: {}{}".format(bcolors.OKGREEN, prev_commit.message, bcolors.ENDC)
+            print(commitStr)
+            print(printableDiff)
+        else:
+            dateStr = "{}Date: {}{}".format(bcolors.OKGREEN, commit_time, bcolors.ENDC)
+            print(dateStr)
+            branchStr = "{}Branch: {}{}".format(bcolors.OKGREEN, branch_name.encode('utf-8'), bcolors.ENDC)
+            print(branchStr)
+            commitStr = "{}Commit: {}{}".format(bcolors.OKGREEN, prev_commit.message.encode('utf-8'), bcolors.ENDC)
+            print(commitStr)
+            print(printableDiff.encode('utf-8'))
+
+def find_strings(git_url, printJson=False):
     output = {"entropicDiffs": []}
+    project_path = clone_git_repo(git_url)
     repo = Repo(project_path)
     already_searched = set()
 
     for remote_branch in repo.remotes.origin.fetch():
-        branch_name = str(remote_branch).split('/')[1]
+        branch_name = remote_branch.name.split('/')[1]
         try:
             repo.git.checkout(remote_branch, b=branch_name)
         except:
@@ -127,14 +152,8 @@ def find_strings(git_url, printJson=False):
                         entropicDiff['diff'] = blob.diff.decode('utf-8', errors='replace') 
                         entropicDiff['stringsFound'] = stringsFound
                         output["entropicDiffs"].append(entropicDiff)
-                        if printJson:
-                            print(json.dumps(output, sort_keys=True, indent=4))
-                        else:
-                            print(bcolors.OKGREEN + "Date: " + commit_time + bcolors.ENDC)
-                            print(bcolors.OKGREEN + "Branch: " + branch_name + bcolors.ENDC)
-                            print(bcolors.OKGREEN + "Commit: " + prev_commit.message + bcolors.ENDC)
-                            print(printableDiff)
 
+                        print_results(printJson, commit_time, branch_name, prev_commit, printableDiff)
             prev_commit = curr_commit
     output["project_path"] = project_path
     return output