From 4a6b540704a138867dca8eb7ff127233de466113 Mon Sep 17 00:00:00 2001 From: Devoalda Date: Mon, 23 Oct 2023 16:52:40 +0800 Subject: [PATCH] Add Virus Scan GRPC --- safeshare/safeshare_app/utils/__init__.py | 0 .../safeshare_app/utils/virusRPC/__init__.py | 0 .../safeshare_app/utils/virusRPC/client.py | 20 ++++++ .../safeshare_app/utils/virusRPC/scan.proto | 24 +++++++ .../safeshare_app/utils/virusRPC/scan_pb2.py | 29 ++++++++ .../safeshare_app/utils/virusRPC/scan_pb2.pyi | 35 ++++++++++ .../utils/virusRPC/scan_pb2_grpc.py | 66 +++++++++++++++++++ .../safeshare_app/utils/virusRPC/server.py | 61 +++++++++++++++++ 8 files changed, 235 insertions(+) create mode 100644 safeshare/safeshare_app/utils/__init__.py create mode 100644 safeshare/safeshare_app/utils/virusRPC/__init__.py create mode 100644 safeshare/safeshare_app/utils/virusRPC/client.py create mode 100644 safeshare/safeshare_app/utils/virusRPC/scan.proto create mode 100644 safeshare/safeshare_app/utils/virusRPC/scan_pb2.py create mode 100644 safeshare/safeshare_app/utils/virusRPC/scan_pb2.pyi create mode 100644 safeshare/safeshare_app/utils/virusRPC/scan_pb2_grpc.py create mode 100644 safeshare/safeshare_app/utils/virusRPC/server.py diff --git a/safeshare/safeshare_app/utils/__init__.py b/safeshare/safeshare_app/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/safeshare/safeshare_app/utils/virusRPC/__init__.py b/safeshare/safeshare_app/utils/virusRPC/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/safeshare/safeshare_app/utils/virusRPC/client.py b/safeshare/safeshare_app/utils/virusRPC/client.py new file mode 100644 index 0000000..98f1043 --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/client.py @@ -0,0 +1,20 @@ +import grpc + +import scan_pb2_grpc +import scan_pb2 + + +class Client: + def __init__(self): + self.channel = grpc.insecure_channel("localhost:50051") + self.stub = scan_pb2_grpc.VirusScanServiceStub(self.channel) + + def ScanFile(self, sha_256_id: str): + response = self.stub.ScanFile(scan_pb2.ScanFileRequest(file_SHA256=sha_256_id)) + print(response) + + +if __name__ == "__main__": + client = Client() + id = "15e4313dddb45875ed67d1ab25f1f5b76f0b3a23e4fa9308c521e3fb30068028" + client.ScanFile(id) diff --git a/safeshare/safeshare_app/utils/virusRPC/scan.proto b/safeshare/safeshare_app/utils/virusRPC/scan.proto new file mode 100644 index 0000000..d002fd7 --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/scan.proto @@ -0,0 +1,24 @@ +syntax = "proto3"; + +package virusscan; + +service VirusScanService { + rpc ScanFile (ScanFileRequest) returns (ScanFileResponse) {} +} + +message ScanFileRequest { + string file_name = 1; + string file_SHA256 = 2; + string file_SHA1 = 3; + string file_MD5 = 4; +} + +message ScanFileResponse { + string file_name = 1; + string file_SHA256 = 2; + string file_SHA1 = 3; + string file_MD5 = 4; + bool is_infected = 5; + string scan_result = 6; + string scan_result_detail = 7; +} diff --git a/safeshare/safeshare_app/utils/virusRPC/scan_pb2.py b/safeshare/safeshare_app/utils/virusRPC/scan_pb2.py new file mode 100644 index 0000000..d42285f --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/scan_pb2.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: scan.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\nscan.proto\x12\tvirusscan\"^\n\x0fScanFileRequest\x12\x11\n\tfile_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66ile_SHA256\x18\x02 \x01(\t\x12\x11\n\tfile_SHA1\x18\x03 \x01(\t\x12\x10\n\x08\x66ile_MD5\x18\x04 \x01(\t\"\xa5\x01\n\x10ScanFileResponse\x12\x11\n\tfile_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66ile_SHA256\x18\x02 \x01(\t\x12\x11\n\tfile_SHA1\x18\x03 \x01(\t\x12\x10\n\x08\x66ile_MD5\x18\x04 \x01(\t\x12\x13\n\x0bis_infected\x18\x05 \x01(\x08\x12\x13\n\x0bscan_result\x18\x06 \x01(\t\x12\x1a\n\x12scan_result_detail\x18\x07 \x01(\t2Y\n\x10VirusScanService\x12\x45\n\x08ScanFile\x12\x1a.virusscan.ScanFileRequest\x1a\x1b.virusscan.ScanFileResponse\"\x00\x62\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'scan_pb2', _globals) +if _descriptor._USE_C_DESCRIPTORS == False: + DESCRIPTOR._options = None + _globals['_SCANFILEREQUEST']._serialized_start=25 + _globals['_SCANFILEREQUEST']._serialized_end=119 + _globals['_SCANFILERESPONSE']._serialized_start=122 + _globals['_SCANFILERESPONSE']._serialized_end=287 + _globals['_VIRUSSCANSERVICE']._serialized_start=289 + _globals['_VIRUSSCANSERVICE']._serialized_end=378 +# @@protoc_insertion_point(module_scope) diff --git a/safeshare/safeshare_app/utils/virusRPC/scan_pb2.pyi b/safeshare/safeshare_app/utils/virusRPC/scan_pb2.pyi new file mode 100644 index 0000000..79c8c50 --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/scan_pb2.pyi @@ -0,0 +1,35 @@ +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from typing import ClassVar as _ClassVar, Optional as _Optional + +DESCRIPTOR: _descriptor.FileDescriptor + +class ScanFileRequest(_message.Message): + __slots__ = ["file_name", "file_SHA256", "file_SHA1", "file_MD5"] + FILE_NAME_FIELD_NUMBER: _ClassVar[int] + FILE_SHA256_FIELD_NUMBER: _ClassVar[int] + FILE_SHA1_FIELD_NUMBER: _ClassVar[int] + FILE_MD5_FIELD_NUMBER: _ClassVar[int] + file_name: str + file_SHA256: str + file_SHA1: str + file_MD5: str + def __init__(self, file_name: _Optional[str] = ..., file_SHA256: _Optional[str] = ..., file_SHA1: _Optional[str] = ..., file_MD5: _Optional[str] = ...) -> None: ... + +class ScanFileResponse(_message.Message): + __slots__ = ["file_name", "file_SHA256", "file_SHA1", "file_MD5", "is_infected", "scan_result", "scan_result_detail"] + FILE_NAME_FIELD_NUMBER: _ClassVar[int] + FILE_SHA256_FIELD_NUMBER: _ClassVar[int] + FILE_SHA1_FIELD_NUMBER: _ClassVar[int] + FILE_MD5_FIELD_NUMBER: _ClassVar[int] + IS_INFECTED_FIELD_NUMBER: _ClassVar[int] + SCAN_RESULT_FIELD_NUMBER: _ClassVar[int] + SCAN_RESULT_DETAIL_FIELD_NUMBER: _ClassVar[int] + file_name: str + file_SHA256: str + file_SHA1: str + file_MD5: str + is_infected: bool + scan_result: str + scan_result_detail: str + def __init__(self, file_name: _Optional[str] = ..., file_SHA256: _Optional[str] = ..., file_SHA1: _Optional[str] = ..., file_MD5: _Optional[str] = ..., is_infected: bool = ..., scan_result: _Optional[str] = ..., scan_result_detail: _Optional[str] = ...) -> None: ... diff --git a/safeshare/safeshare_app/utils/virusRPC/scan_pb2_grpc.py b/safeshare/safeshare_app/utils/virusRPC/scan_pb2_grpc.py new file mode 100644 index 0000000..967a39e --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/scan_pb2_grpc.py @@ -0,0 +1,66 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +import scan_pb2 as scan__pb2 + + +class VirusScanServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.ScanFile = channel.unary_unary( + '/virusscan.VirusScanService/ScanFile', + request_serializer=scan__pb2.ScanFileRequest.SerializeToString, + response_deserializer=scan__pb2.ScanFileResponse.FromString, + ) + + +class VirusScanServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def ScanFile(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_VirusScanServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'ScanFile': grpc.unary_unary_rpc_method_handler( + servicer.ScanFile, + request_deserializer=scan__pb2.ScanFileRequest.FromString, + response_serializer=scan__pb2.ScanFileResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'virusscan.VirusScanService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class VirusScanService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def ScanFile(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, '/virusscan.VirusScanService/ScanFile', + scan__pb2.ScanFileRequest.SerializeToString, + scan__pb2.ScanFileResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/safeshare/safeshare_app/utils/virusRPC/server.py b/safeshare/safeshare_app/utils/virusRPC/server.py new file mode 100644 index 0000000..c7afc7e --- /dev/null +++ b/safeshare/safeshare_app/utils/virusRPC/server.py @@ -0,0 +1,61 @@ +import grpc +import scan_pb2 +import scan_pb2_grpc +from concurrent import futures +import requests + +apiKey = "" # API key from VirusTotal + +headers = { + "accept": "application/json", + "x-apikey": apiKey +} + + +class VirusScanServicer(scan_pb2_grpc.VirusScanServiceServicer): + url = "https://www.virustotal.com/api/v3/files/" + + def ScanFile(self, request, context): + result = self.ScanSHA(request.file_SHA256) + return result + + def ScanSHA(self, sha256_hash: str) -> scan_pb2.ScanFileResponse: + self.url += sha256_hash + + response = requests.get(self.url, headers=headers) + data = response.json()["data"] + result = scan_pb2.ScanFileResponse() + + result.is_infected = data["attributes"]["last_analysis_stats"]["malicious"] > 0 + + result.file_name = data["attributes"]["names"][0] + result.file_SHA1 = data["attributes"]["sha1"] + result.file_SHA256 = data["attributes"]["sha256"] + + return result + + +class VirusScanServer: + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + scan_pb2_grpc.add_VirusScanServiceServicer_to_server(VirusScanServicer(), server) + server.add_insecure_port("[::]:50051") + + def serve(self): + try: + self.server.start() + print("Server started, listening on 50051") + self.server.wait_for_termination() + except KeyboardInterrupt: + print("Server stopped") + self.server.stop(0) + + def __del__(self): + self.server.stop(0) + + def __enter__(self): + return self + + +if __name__ == "__main__": + cal = VirusScanServer() + cal.serve()