diff --git a/ACMAS/.env b/ACMAS/.env index 16c2c81..0a4143e 100644 --- a/ACMAS/.env +++ b/ACMAS/.env @@ -16,3 +16,4 @@ SECRET_KEY=SECRET_KEY SQL_ENGINE=django.db.backends.postgresql SQL_HOST=db SQL_PORT=5432 +VT_API_KEY=INSERT_KEY_HERE diff --git a/ACMAS/app/ACMAS_Web/file_scanning.py b/ACMAS/app/ACMAS_Web/file_scanning.py new file mode 100644 index 0000000..f4ae4eb --- /dev/null +++ b/ACMAS/app/ACMAS_Web/file_scanning.py @@ -0,0 +1,41 @@ +import hashlib +import os +import vt + + +def hash_file(file_path): + # make a hash object + h = hashlib.sha1() + + # open file for reading in binary mode + with open(file_path,'rb') as file: + # loop till the end of the file + chunk = 0 + while chunk != b'': + # read only 1024 bytes at a time + chunk = file.read(1024) + h.update(chunk) + + # return the hex representation of digest + return h.hexdigest() + +def scan_file(file_path): + VT_API_KEY = os.environ.get("VT_API_KEY", default="") + client = vt.Client(apikey=VT_API_KEY, agent='ACMAS') + # file_hash = hash_file(file_path) + file = client.get_object(file_hash) + if file.last_analysis_stats['malicious'] > 0: + return 'malicious' + elif file.last_analysis_stats['suspicious'] > 0: + return 'suspicious' + elif file.last_analysis_stats['harmless'] > 0: + with open(file_path, 'rb') as f: + # Scan the file and block until it is scanned + analysis = client.scan_file(f, wait_for_completion=True) + + # try: + # analysis = client.scan_file(file_path, wait_for_completion=True) + # return analysis + # except vt.APIError as e: + # print(e) + # return None diff --git a/ACMAS/app/ACMAS_Web/upload.py b/ACMAS/app/ACMAS_Web/upload.py index ad191ce..154c231 100644 --- a/ACMAS/app/ACMAS_Web/upload.py +++ b/ACMAS/app/ACMAS_Web/upload.py @@ -1,5 +1,6 @@ import zlib from datetime import date +from file_scanning import scan_file from django.core.files.base import ContentFile from django.core.files.storage import FileSystemStorage @@ -59,6 +60,8 @@ def uploadFile(self, uni, course, fType, file): file_url = fs.url(savedFile) # Retrieve the file path print(f'FILE "{savedFile}" uploaded to "{file_url}"\n') + scan_result = scan_file(file_url) + # Adding file to database db_file = UploadedFile( filename=savedFile, diff --git a/ACMAS/app/requirements.txt b/ACMAS/app/requirements.txt index 98cf453..5ae5d00 100644 --- a/ACMAS/app/requirements.txt +++ b/ACMAS/app/requirements.txt @@ -3,4 +3,5 @@ django-dbbackup==4.0.2 django_components psycopg2-binary==2.9.7 -gunicorn==20.1.0 \ No newline at end of file +gunicorn==20.1.0 +vt-py==0.17.5 \ No newline at end of file