79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import re
|
|
import grpc
|
|
import dynamo_pb2 as pb2
|
|
import dynamo_pb2_grpc as pb2_grpc
|
|
|
|
import boto3 as boto # 1.28.68
|
|
|
|
# dynamo db instance
|
|
dynamodb = boto.resource('dynamodb')
|
|
sha256_table = dynamodb.Table('safeshare_sha256')
|
|
sha1_table = dynamodb.Table('safeshare_sha1')
|
|
md5_table = dynamodb.Table('safeshare_md5')
|
|
|
|
# hash
|
|
hex_pattern = re.compile("^[a-fA-F0-9]+$")
|
|
|
|
|
|
def check_sha256(sha256):
|
|
response = sha256_table.get_item(Key={'sha256': sha256})
|
|
return 'Item' in response
|
|
|
|
|
|
def check_sha1(sha1):
|
|
response = sha1_table.get_item(Key={'sha1': sha1})
|
|
return 'Item' in response
|
|
|
|
|
|
def check_md5(md5):
|
|
response = md5_table.get_item(Key={'md5': md5})
|
|
return 'Item' in response
|
|
|
|
|
|
class Dynamo(pb2_grpc.Dynamo_DBServicer):
|
|
def CheckHash(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:
|
|
return pb2.Response(is_exist=check_sha256(request.file_hash))
|
|
elif length == 40:
|
|
return pb2.Response(is_exist=check_sha1(request.file_hash))
|
|
elif length == 32:
|
|
return pb2.Response(is_exist=check_md5(request.file_hash))
|
|
else:
|
|
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)
|
|
|
|
|
|
def serve():
|
|
server = grpc.server(ThreadPoolExecutor(max_workers=10))
|
|
pb2_grpc.add_Dynamo_DBServicer_to_server(Dynamo(), server)
|
|
|
|
server.add_insecure_port('[::]:50051')
|
|
server.start()
|
|
server.wait_for_termination()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
serve()
|