From cad048119d8493327a0b932c6404e250018d07b4 Mon Sep 17 00:00:00 2001
From: Oran Nadler <orann>
Date: Tue, 15 May 2018 15:26:46 +0300
Subject: [PATCH] split pth to be able to run standalone and in website

---
 monkey_island/cc/resources/pthmap.py    | 21 +++++++++++++
 monkey_island/cc/resources/pthreport.py | 41 ++++++-------------------
 2 files changed, 30 insertions(+), 32 deletions(-)
 create mode 100644 monkey_island/cc/resources/pthmap.py

diff --git a/monkey_island/cc/resources/pthmap.py b/monkey_island/cc/resources/pthmap.py
new file mode 100644
index 000000000..c59bb90cd
--- /dev/null
+++ b/monkey_island/cc/resources/pthmap.py
@@ -0,0 +1,21 @@
+import hashlib
+import binascii
+import copy
+import flask_restful
+from pthreport import PassTheHashReport
+
+from cc.auth import jwt_required
+from cc.services.edge import EdgeService
+from cc.services.node import NodeService
+from cc.database import mongo
+
+class PthMap(flask_restful.Resource):
+    @jwt_required()
+    def get(self, **kw):
+        graph = PassTheHashReport()
+
+        return \
+            {
+                "nodes": [{"id": x, "label": Machine(x).GetIp()} for x in graph.vertices],
+                "edges": [{"id": str(s) + str(t), "from": s, "to": t, "label": label} for s, t, label in graph.edges]
+            }
diff --git a/monkey_island/cc/resources/pthreport.py b/monkey_island/cc/resources/pthreport.py
index 7472a4476..664732dfb 100644
--- a/monkey_island/cc/resources/pthreport.py
+++ b/monkey_island/cc/resources/pthreport.py
@@ -1,38 +1,15 @@
 import hashlib
 import binascii
 import copy
-from pymongo import MongoClient
 
-class mongo(object):
-    db = MongoClient().monkeyisland
+if __name__ == "__main__":
+    from pymongo import MongoClient
 
-#class PthMap(flask_restful.Resource):
-class PthMap(object):
-    #@jwt_required()
-    def get(self, **kw):
-        graph = PassTheHashMap()
-
-        return \
-            {
-                "nodes": [{"id": x, "label": Machine(x).GetIp()} for x in graph.vertices],
-                "edges": [{"id": str(s) + str(t), "from": s, "to": t, "label": label} for s, t, label in graph.edges]
-            }
-
-if not __name__ == "__main__":
-    import flask_restful
-    
-    from cc.auth import jwt_required
-    from cc.services.edge import EdgeService
-    from cc.services.node import NodeService
+    class mongo(object):
+        db = MongoClient().monkeyisland
+else:
     from cc.database import mongo
     
-    PthMapOrig = PthMap
-    
-    class PthMap(flask_restful.Resource):
-        @jwt_required()
-        def get(self, **kw):
-            return PthMapOrig.get(self, **kw)
-    
 DsRole_RoleStandaloneWorkstation = 0
 DsRole_RoleMemberWorkstation = 1
 DsRole_RoleStandaloneServer = 2
@@ -88,8 +65,8 @@ def cache(foo):
         elif type(o) == PthMap:
             return "PthMapSingleton"
             
-        elif type(o) == PassTheHashMap:
-            return "PassTheHashMapSingleton"
+        elif type(o) == PassTheHashReport:
+            return "PassTheHashReportSingleton"
 
         else:
             assert False, "%s of type %s is not hashable" % (repr(o), type(o))
@@ -564,7 +541,7 @@ class Machine(object):
         
         return names
 
-class PassTheHashMap(object):
+class PassTheHashReport(object):
     def __init__(self):
         self.vertices = self.GetAllMachines()
         
@@ -944,7 +921,7 @@ class PassTheHashMap(object):
         return shared_admins
 
 def main():
-    pth = PassTheHashMap()
+    pth = PassTheHashReport()
 
     print """<style>table, td {border: ridge;} .hidden * {color:white; border:white;}</style>"""
     print "<h1>Pass The Hash Report</h1>"