import vdb client

This commit is contained in:
Richie 2023-10-29 22:31:57 +08:00
parent a1e10f2474
commit c38ded617d
7 changed files with 177 additions and 2 deletions

View File

@ -0,0 +1,24 @@
import grpc
import dynamo_pb2 as pb2
import dynamo_pb2_grpc as pb2_grpc
class Client:
def __init__(self):
self.channel = grpc.insecure_channel("localhost:50051")
self.stub = pb2_grpc.Dynamo_DBStub(self.channel)
def CheckFile(self, sha_256_id: str):
response = self.stub.CheckHash(pb2.Request(file_hash=sha_256_id))
return response.is_exist
def UpdateFile(self, sha_256_id: str):
response = self.stub.UpdateHash(pb2.Request(file_hash=sha_256_id))
return response.is_exist
# if __name__ == "__main__":
# client = Client()
# id = "15e4313dddb45875ed67d1ab25f1f5b76f0b3a23e4fa9308c521e3fb30068028"
# print(client.CheckFile(id))
# # client.UpdateFile(id)

View File

@ -0,0 +1,16 @@
syntax = "proto3";
package dynamo_db;
service Dynamo_DB {
rpc CheckHash (Request) returns (Response) {}
rpc UpdateHash (Request) returns (Response) {}
}
message Request {
string file_hash = 2;
}
message Response {
bool is_exist = 1;
}

View File

@ -0,0 +1,29 @@
# -*- coding: utf-8 -*-
# Generated by the protocol buffer compiler. DO NOT EDIT!
# source: dynamo.proto
"""Generated protocol buffer code."""
from google.protobuf import descriptor as _descriptor
from google.protobuf import descriptor_pool as _descriptor_pool
from google.protobuf import symbol_database as _symbol_database
from google.protobuf.internal import builder as _builder
# @@protoc_insertion_point(imports)
_sym_db = _symbol_database.Default()
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0c\x64ynamo.proto\x12\tdynamo_db\"\x1c\n\x07Request\x12\x11\n\tfile_hash\x18\x02 \x01(\t\"\x1c\n\x08Response\x12\x10\n\x08is_exist\x18\x01 \x01(\x08\x32|\n\tDynamo_DB\x12\x36\n\tCheckHash\x12\x12.dynamo_db.Request\x1a\x13.dynamo_db.Response\"\x00\x12\x37\n\nUpdateHash\x12\x12.dynamo_db.Request\x1a\x13.dynamo_db.Response\"\x00\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'dynamo_pb2', _globals)
if _descriptor._USE_C_DESCRIPTORS == False:
DESCRIPTOR._options = None
_globals['_REQUEST']._serialized_start=27
_globals['_REQUEST']._serialized_end=55
_globals['_RESPONSE']._serialized_start=57
_globals['_RESPONSE']._serialized_end=85
_globals['_DYNAMO_DB']._serialized_start=87
_globals['_DYNAMO_DB']._serialized_end=211
# @@protoc_insertion_point(module_scope)

View File

@ -0,0 +1,99 @@
# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
"""Client and server classes corresponding to protobuf-defined services."""
import grpc
import dynamo_pb2 as dynamo__pb2
class Dynamo_DBStub(object):
"""Missing associated documentation comment in .proto file."""
def __init__(self, channel):
"""Constructor.
Args:
channel: A grpc.Channel.
"""
self.CheckHash = channel.unary_unary(
'/dynamo_db.Dynamo_DB/CheckHash',
request_serializer=dynamo__pb2.Request.SerializeToString,
response_deserializer=dynamo__pb2.Response.FromString,
)
self.UpdateHash = channel.unary_unary(
'/dynamo_db.Dynamo_DB/UpdateHash',
request_serializer=dynamo__pb2.Request.SerializeToString,
response_deserializer=dynamo__pb2.Response.FromString,
)
class Dynamo_DBServicer(object):
"""Missing associated documentation comment in .proto file."""
def CheckHash(self, request, context):
"""Missing associated documentation comment in .proto file."""
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
def UpdateHash(self, request, context):
"""Missing associated documentation comment in .proto file."""
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
def add_Dynamo_DBServicer_to_server(servicer, server):
rpc_method_handlers = {
'CheckHash': grpc.unary_unary_rpc_method_handler(
servicer.CheckHash,
request_deserializer=dynamo__pb2.Request.FromString,
response_serializer=dynamo__pb2.Response.SerializeToString,
),
'UpdateHash': grpc.unary_unary_rpc_method_handler(
servicer.UpdateHash,
request_deserializer=dynamo__pb2.Request.FromString,
response_serializer=dynamo__pb2.Response.SerializeToString,
),
}
generic_handler = grpc.method_handlers_generic_handler(
'dynamo_db.Dynamo_DB', rpc_method_handlers)
server.add_generic_rpc_handlers((generic_handler,))
# This class is part of an EXPERIMENTAL API.
class Dynamo_DB(object):
"""Missing associated documentation comment in .proto file."""
@staticmethod
def CheckHash(request,
target,
options=(),
channel_credentials=None,
call_credentials=None,
insecure=False,
compression=None,
wait_for_ready=None,
timeout=None,
metadata=None):
return grpc.experimental.unary_unary(request, target, '/dynamo_db.Dynamo_DB/CheckHash',
dynamo__pb2.Request.SerializeToString,
dynamo__pb2.Response.FromString,
options, channel_credentials,
insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
@staticmethod
def UpdateHash(request,
target,
options=(),
channel_credentials=None,
call_credentials=None,
insecure=False,
compression=None,
wait_for_ready=None,
timeout=None,
metadata=None):
return grpc.experimental.unary_unary(request, target, '/dynamo_db.Dynamo_DB/UpdateHash',
dynamo__pb2.Request.SerializeToString,
dynamo__pb2.Response.FromString,
options, channel_credentials,
insecure, call_credentials, compression, wait_for_ready, timeout, metadata)

View File

@ -11,6 +11,11 @@ from django.http import HttpResponse
from rest_framework.response import Response from rest_framework.response import Response
from rest_framework.views import APIView from rest_framework.views import APIView
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../utils/safeshare_vdb_client")
import client
class ManageItemsView(APIView): class ManageItemsView(APIView):
def post(self, request): def post(self, request):
@ -52,8 +57,8 @@ class ManageItemsView(APIView):
# If RPC client import fails, skip virus scan # If RPC client import fails, skip virus scan
# Call RPC For virus scan # Call RPC For virus scan
try: try:
client = Client() grpc_client = client.Client()
result = client.CheckFile(hash_signature) result = grpc_client.CheckFile(hash_signature)
except Exception as e: except Exception as e:
result = False result = False

View File

@ -2,6 +2,7 @@ import grpc
import dynamo_pb2 as pb2 import dynamo_pb2 as pb2
import dynamo_pb2_grpc as pb2_grpc import dynamo_pb2_grpc as pb2_grpc
class Client: class Client:
def __init__(self): def __init__(self):
self.channel = grpc.insecure_channel("localhost:50051") self.channel = grpc.insecure_channel("localhost:50051")

View File

@ -87,6 +87,7 @@ class Dynamo(pb2_grpc.Dynamo_DBServicer):
else: else:
length = len(request.file_hash) length = len(request.file_hash)
if length == 64: if length == 64:
print("check sha256")
return pb2.Response(is_exist=check_sha256(request.file_hash)) return pb2.Response(is_exist=check_sha256(request.file_hash))
elif length == 40: elif length == 40:
return pb2.Response(is_exist=check_sha1(request.file_hash)) return pb2.Response(is_exist=check_sha1(request.file_hash))