integrate totalvirus api

safeshare/safeshare_vdb/client.py

@@ -8,11 +8,11 @@ class Client:
         self.stub = pb2_grpc.Dynamo_DBStub(self.channel)
 
     def CheckFile(self, sha_256_id: str):
-        response = self.stub.CheckFile(pb2.Request(file_hash=sha_256_id))
+        response = self.stub.CheckHash(pb2.Request(file_hash=sha_256_id))
         print(response)
 
     def UpdateFile(self, sha_256_id: str):
-        response = self.stub.UpdateFile(pb2.Request(file_hash=sha_256_id))
+        response = self.stub.UpdateHash(pb2.Request(file_hash=sha_256_id))
         print(response)
 
 
@@ -20,4 +20,4 @@ if __name__ == "__main__":
     client = Client()
     id = "15e4313dddb45875ed67d1ab25f1f5b76f0b3a23e4fa9308c521e3fb30068028"
     client.CheckFile(id)
-    client.UpdateFile(id)
+    # client.UpdateFile(id)

safeshare/safeshare_vdb/env

@@ -0,0 +1 @@
+API_TOKEN=

safeshare/safeshare_vdb/server.py

@@ -4,6 +4,9 @@ import re
 import grpc
 import dynamo_pb2 as pb2
 import dynamo_pb2_grpc as pb2_grpc
+import environ
+import os
+import requests
 
 import boto3 as boto  # 1.28.68
 
@@ -13,23 +16,68 @@ sha256_table = dynamodb.Table('safeshare_sha256')
 sha1_table = dynamodb.Table('safeshare_sha1')
 md5_table = dynamodb.Table('safeshare_md5')
 
+# TotalVirus API key
+environ.Env.read_env('./.env')
+api = environ.Env().str('API_TOKEN')
+
+headers = {
+    "accept": "application/json",
+    "x-apikey": api
+}
+
 # hash
 hex_pattern = re.compile("^[a-fA-F0-9]+$")
 
 
+def upload(hash_val):
+    if not hex_pattern.match(hash_val):
+        return False
+    else:
+        length = len(hash_val)
+        if length == 64:
+            sha256_table.put_item(Item={'sha256': hash_val})
+        elif length == 40:
+            sha1_table.put_item(Item={'sha1': hash_val})
+        elif length == 32:
+            md5_table.put_item(Item={'md5': hash_val})
+        else:
+            return False
+
+        return True
+
+
+def scan(hash_val):
+    url = "https://www.virustotal.com/api/v3/files/" + hash_val
+    response = requests.get(url, headers=headers)
+    if response.status_code == 200:
+        data = response.json()["data"]
+        return data["attributes"]["last_analysis_stats"]["malicious"] > 0
+    else:
+        return False
+
+
 def check_sha256(sha256):
     response = sha256_table.get_item(Key={'sha256': sha256})
-    return 'Item' in response
+    if "Item" not in response:
+        return upload(sha256) if scan(sha256) else False
+    else:
+        return True
 
 
 def check_sha1(sha1):
     response = sha1_table.get_item(Key={'sha1': sha1})
-    return 'Item' in response
+    if "Item" not in response:
+        return upload(sha1) if scan(sha1) else False
+    else:
+        return True
 
 
 def check_md5(md5):
     response = md5_table.get_item(Key={'md5': md5})
-    return 'Item' in response
+    if "Item" not in response:
+        return upload(md5) if scan(md5) else False
+    else:
+        return True
 
 
 class Dynamo(pb2_grpc.Dynamo_DBServicer):
@@ -48,21 +96,7 @@ class Dynamo(pb2_grpc.Dynamo_DBServicer):
                 return pb2.Response(is_exist=False)
 
     def UpdateHash(self, request, context):
-        if not hex_pattern.match(request.file_hash):
-            return pb2.Response(is_exist=False)
-        else:
-            length = len(request.file_hash)
-            if length == 64:
-                sha256_table.put_item(Item={'sha256': request.file_hash})
-                return pb2.Response(is_exist=True)
-            elif length == 40:
-                sha1_table.put_item(Item={'sha1': request.file_hash})
-                return pb2.Response(is_exist=True)
-            elif length == 32:
-                md5_table.put_item(Item={'md5': request.file_hash})
-                return pb2.Response(is_exist=True)
-            else:
-                return pb2.Response(is_exist=False)
+        return pb2.Response(is_exist=upload(request.file_hash))
 
 
 def serve():