diff --git a/safeshare/safeshare_vdb/__init__.py b/safeshare/safeshare_vdb/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/safeshare/safeshare_vdb/client.py b/safeshare/safeshare_vdb/client.py index 0c3b5f5..8e42fee 100644 --- a/safeshare/safeshare_vdb/client.py +++ b/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) diff --git a/safeshare/safeshare_vdb/env b/safeshare/safeshare_vdb/env new file mode 100644 index 0000000..c65e6b3 --- /dev/null +++ b/safeshare/safeshare_vdb/env @@ -0,0 +1 @@ +API_TOKEN= \ No newline at end of file diff --git a/safeshare/safeshare_vdb/server.py b/safeshare/safeshare_vdb/server.py index e8a875a..bee491e 100644 --- a/safeshare/safeshare_vdb/server.py +++ b/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():