From 989dcc578f70fdb932c7cabc6892855782bdfed7 Mon Sep 17 00:00:00 2001
From: Richie <2837357W@student.gla.ac.uk>
Date: Mon, 30 Oct 2023 11:35:07 +0800
Subject: [PATCH] safeshare_vdb dockerfile

---
 safeshare/safeshare_vdb/Dockerfile |  5 ++++-
 safeshare/safeshare_vdb/server.py  | 18 ++++++++++++------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/safeshare/safeshare_vdb/Dockerfile b/safeshare/safeshare_vdb/Dockerfile
index 1c1e57f..ccc6e78 100644
--- a/safeshare/safeshare_vdb/Dockerfile
+++ b/safeshare/safeshare_vdb/Dockerfile
@@ -8,7 +8,10 @@ COPY . /app
 WORKDIR /app
 
 # Install any needed packages for grpc and environ
-RUN pip install grpcio grpcio-tools django-environ requests boto3
+RUN pip install grpcio grpcio-tools requests boto3 python-environ protobuf
+
+# Make port 50051 available to the world outside this container
+EXPOSE 50051
 
 # Start Server at server.py
 CMD ["python", "server.py"]
\ No newline at end of file
diff --git a/safeshare/safeshare_vdb/server.py b/safeshare/safeshare_vdb/server.py
index 32cec6a..3bf68be 100644
--- a/safeshare/safeshare_vdb/server.py
+++ b/safeshare/safeshare_vdb/server.py
@@ -10,16 +10,22 @@ import requests
 
 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')
-
 # TotalVirus API key
 environ.Env.read_env('./.env')
 api = environ.Env().str('API_TOKEN')
 
+# dynamo db instance
+session = boto.Session(
+    aws_access_key_id=os.environ.get('AWS_ACCESS_KEY_ID'),
+    aws_secret_access_key=os.environ.get('AWS_SECRET_ACCESS_KEY'),
+    region_name=os.environ.get('REGION')
+)
+
+dynamodb = session.resource('dynamodb')
+sha256_table = dynamodb.Table('safeshare_sha256')
+sha1_table = dynamodb.Table('safeshare_sha1')
+md5_table = dynamodb.Table('safeshare_md5')
+
 headers = {
     "accept": "application/json",
     "x-apikey": api