Merge pull request #1918 from raft-tech/feat/1133-files-transferred-t…

…o-ACF-Titan

feat 1133 files transfer to ACF Titan

tdrs-backend/Pipfile

@@ -20,7 +20,8 @@ pytest-mock = "*"
 pytest-factoryboy = "*"
 
 [packages]
-boto3 = "==1.17.112"
+pytz = "2022.1"
+boto3 = "==1.24.51"
 cryptography = "==3.4.7"
 dj-database-url = "==0.5.0"
 django = "==3.2.13"
@@ -45,6 +46,12 @@ pyjwt = "==2.4.0"
 requests = "==2.27.1"
 wait-for-it = "==2.2.0"
 requests-mock = "==1.9.3"
+celery = "==5.2.6"
+redis = "==4.1.2"
+flower = "==1.1.0"
+django-celery-beat = "==2.2.1"
+paramiko = "==2.11.0"
+pytest_sftpserver = "==1.3.0"
 
 [requires]
 python_version = "3.10.4"

tdrs-backend/apt.yml

@@ -5,3 +5,5 @@ repos:
   - deb http://apt.postgresql.org/pub/repos/apt/ bionic-pgdg main
 packages:
   - postgresql-client-12
+  - libjemalloc-dev
+  - redis

tdrs-backend/docker-compose.local.yml

@@ -22,6 +22,8 @@ services:
       - DATA_DIR=/tmp/localstack/data
       - AWS_BUCKET=tdp-datafiles-localstack
       - AWS_REGION_NAME=us-gov-west-1
+      - HOSTNAME=localstack
+      - HOSTNAME_EXTERNAL=localstack
     ports:
       - "4566:4566"
     volumes:

tdrs-backend/docker-compose.yml

@@ -67,19 +67,35 @@ services:
       - AMS_CLIENT_ID
       - AMS_CLIENT_SECRET
       - AMS_CONFIGURATION_ENDPOINT
+      - ACFTITAN_HOST
+      - ACFTITAN_KEY
+      - ACFTITAN_USERNAME
+      - REDIS_URI
+      - REDIS_SERVER_LOCAL=TRUE
+      - ACFTITAN_SFTP_PYTEST
     volumes:
       - .:/tdpapp
     image: tdp
     build: .
     command: >
       bash -c "./wait_for_services.sh &&
-               ./gunicorn_start.sh"
+               ./gunicorn_start.sh && celery -A tdpservice.settings worker -l info"
     ports:
       - "8080:8080"
+      - "5555:5555"
     depends_on:
       - clamav-rest
       - localstack
       - postgres
+      - redis-server
+
+  redis-server:
+    image: "redis:alpine"
+    command: redis-server /tdpapp/redis.conf
+    ports:
+      - "6379:6379"
+    volumes:
+      - .:/tdpapp
 
 volumes:
   localstack_data:

tdrs-backend/gunicorn_start.sh

@@ -1,11 +1,30 @@
 #!/usr/bin/env bash
 # Apply database migrations
 set -e
+
+echo "REDIS_SERVER"
+echo $REDIS_SERVER_LOCAL
+if [[ "$REDIS_SERVER_LOCAL" = "TRUE" || "$CIRCLE_JOB" = "backend-owasp-scan" ]]; then
+    echo "Run redis server on docker"
+else
+    echo "Run redis server locally"
+    export LD_LIBRARY_PATH=/home/vcap/deps/0/lib/:/home/vcap/deps/1/lib:$LD_LIBRARY_PATH
+    ( cd  /home/vcap/deps/0/bin/; ./redis-server /home/vcap/app/redis.conf &)
+fi
+
+#
 echo "Applying database migrations"
 python manage.py makemigrations
 python manage.py migrate
 python manage.py populate_stts
 python manage.py collectstatic --noinput
+
+celery -A tdpservice.settings worker -c 1 --max-memory-per-child 5000 &
+sleep 5
+# TODO: Uncomment the following line to add flower service when memory limitation is resolved
+# celery -A tdpservice.settings --broker=$REDIS_URI flower &
+celery -A tdpservice.settings beat -l info --scheduler django_celery_beat.schedulers:DatabaseScheduler &
+
 echo "Starting Gunicorn"
 if [[ "$DJANGO_CONFIGURATION" = "Development" || "$DJANGO_CONFIGURATION" = "Local" ]]; then
     gunicorn_params="--bind 0.0.0.0:8080 --timeout 10 --workers 3 --reload --log-level $LOGGING_LEVEL"
@@ -15,4 +34,4 @@ fi
 
 gunicorn_cmd="gunicorn tdpservice.wsgi:application $gunicorn_params"
 
-exec $gunicorn_cmd
+exec $gunicorn_cmd

tdrs-backend/redis.conf

@@ -0,0 +1,7 @@
+#############NETWORK################
+bind 0.0.0.0
+port 6379
+tcp-keepalive 60
+##############LIMIT MEMORY#############
+maxmemory 20mb
+maxmemory-policy volatile-lru

tdrs-backend/scripts/sshkey_change.py

@@ -0,0 +1,12 @@
+"""Script to change EOL with _ which is required for the key to be used as env variable."""
+
+import sys
+
+if __name__ == "__main__":
+    key_file_name = sys.argv[1]
+    with open(key_file_name, 'r') as key_file:
+        key_string = key_file.read()
+        key_string = key_string.replace('\n', '_')
+        with open(key_file_name + '_', 'w') as key_file:
+            key_file.write(key_string)
+            print('The key is saved in: ', key_file_name + '_')

tdrs-backend/tdpservice/data_files/models.py

@@ -149,10 +149,16 @@ class Meta:
     @property
     def filename(self):
         """Return the correct filename for this data file."""
-        return self.stt.filenames.get(self.section, self.create_filename())
+        # TODO: This is interim logic, it has to be changed when all sections are available to requester
+        if str(self.stt.type).lower() == 'tribe':
+            return self.stt.filenames.get(('Tribal ' if 'Tribal' not in self.section else '') + self.section,
+                                          self.create_filename())
+        else:
+            return self.stt.filenames.get(self.section, self.create_filename())
 
     def create_filename(self, prefix='ADS.E2J'):
         """Return a valid file name for sftp transfer."""
+        """TODO: This method has to be removed"""
         # STT_TYPES = ["state", "territory", "tribe"]
         SECTION = [i.value for i in list(self.Section)]
 

tdrs-backend/tdpservice/data_files/serializers.py

@@ -1,5 +1,5 @@
 """Serialize stt data."""
-
+import logging
 from rest_framework import serializers
 
 from tdpservice.data_files.errors import ImmutabilityError
@@ -11,7 +11,7 @@
 from tdpservice.security.models import ClamAVFileScan
 from tdpservice.stts.models import STT
 from tdpservice.users.models import User
-
+logger = logging.getLogger(__name__)
 
 class DataFileSerializer(serializers.ModelSerializer):
     """Serializer for Data files."""

tdrs-backend/tdpservice/data_files/test/test_models.py

@@ -3,7 +3,7 @@
 
 from tdpservice.stts.models import STT
 
-from ..models import DataFile
+from tdpservice.data_files.models import DataFile
 
 
 @pytest.mark.django_db
@@ -81,4 +81,7 @@ def test_data_files_filename_is_expected(user):
                 "user": user,
                 "stt": stt
             })
-            assert new_data_file.filename == stt.filenames[section]
+            if stt.type == 'tribe':
+                assert new_data_file.filename == stt.filenames['Tribal ' if 'Tribal' not in section else '' + section]
+            else:
+                assert new_data_file.filename == stt.filenames[section]

tdrs-backend/tdpservice/data_files/views.py

@@ -1,8 +1,8 @@
 """Check if user is authorized."""
-import logging
 
 from django.http import StreamingHttpResponse
 from django_filters import rest_framework as filters
+from django.conf import settings
 from drf_yasg.openapi import Parameter
 from drf_yasg.utils import swagger_auto_schema
 from rest_framework.parsers import MultiPartParser
@@ -12,12 +12,12 @@
 from rest_framework.viewsets import ModelViewSet
 from rest_framework.decorators import action
 from wsgiref.util import FileWrapper
+from rest_framework import status
 
 from tdpservice.data_files.serializers import DataFileSerializer
 from tdpservice.data_files.models import DataFile
 from tdpservice.users.permissions import DataFilePermissions
-
-logger = logging.getLogger()
+from tdpservice.scheduling import sftp_task
 
 
 class DataFileFilter(filters.FilterSet):
@@ -51,6 +51,21 @@ class DataFileViewSet(ModelViewSet):
     # we will be able to appropriately refer to the latest versions only.
     ordering = ['-version']
 
+    def create(self, request, *args, **kwargs):
+        """Override create to upload in case of successful scan."""
+        response = super().create(request, *args, **kwargs)
+
+        # Upload to ACF-TITAN only if file is passed the virus scan and created
+        if response.status_code == status.HTTP_201_CREATED or response.status_code == status.HTTP_200_OK:
+            sftp_task.upload.delay(
+                data_file_pk=response.data.get('id'),
+                server_address=settings.ACFTITAN_SERVER_ADDRESS,
+                local_key=settings.ACFTITAN_LOCAL_KEY,
+                username=settings.ACFTITAN_USERNAME,
+                port=22
+            )
+        return response
+
     def filter_queryset(self, queryset):
         """Only apply filters to the list action."""
         if self.action != 'list':

tdrs-backend/tdpservice/scheduling/__init__.py

@@ -0,0 +1 @@
+

tdrs-backend/tdpservice/scheduling/admin.py

@@ -0,0 +1,3 @@
+"""Admin class for Scheduling tasks objects."""
+
+# Register your models here.

tdrs-backend/tdpservice/scheduling/apps.py

@@ -0,0 +1,10 @@
+"""Scheduling app configuration."""
+
+from django.apps import AppConfig
+
+
+class TasksConfig(AppConfig):
+    """Scheduling task config."""
+
+    default_auto_field = 'django.db.models.BigAutoField'
+    name = 'tdpservice.scheduling'

tdrs-backend/tdpservice/scheduling/sftp_task.py

@@ -0,0 +1,111 @@
+"""schedule tasks."""
+from __future__ import absolute_import
+# The tasks
+
+import hashlib
+import os
+
+from celery import shared_task
+from django.conf import settings
+import datetime
+import paramiko
+import logging
+from tdpservice.data_files.models import DataFile, LegacyFileTransfer
+
+logger = logging.getLogger(__name__)
+
+
+@shared_task
+def upload(data_file_pk,
+           server_address=settings.ACFTITAN_SERVER_ADDRESS,
+           local_key=settings.ACFTITAN_LOCAL_KEY,
+           username=settings.ACFTITAN_USERNAME,
+           port=22
+           ):
+    """
+    Upload to SFTP server.
+
+    This task uploads the file in DataFile object with pk = data_file_pk
+    to sftp server as defined in Settings file
+    """
+    # Upload file
+    data_file = DataFile.objects.get(id=data_file_pk)
+    file_transfer_record = LegacyFileTransfer(
+        data_file=data_file,
+        uploaded_by=data_file.user,
+        file_name=data_file.filename,
+    )
+
+    def write_key_to_file(private_key):
+        """Paramiko require the key in file object format."""
+        with open('temp_key_file', 'w') as f:
+            f.write(private_key)
+            f.close()
+        return 'temp_key_file'
+
+    def create_dir(directory_name, sftp_server):
+        """Code snippet to create directory in SFTP server."""
+        try:
+            sftp_server.chdir(directory_name)  # Test if remote_path exists
+        except IOError:
+            sftp_server.mkdir(directory_name)  # Create remote_path
+            sftp_server.chdir(directory_name)
+
+    try:
+        # Create directory names for ACF titan
+        destination = str(data_file.filename)
+        today_date = datetime.datetime.today()
+        upper_directory_name = today_date.strftime('%Y%m%d')
+        lower_directory_name = today_date.strftime(str(data_file.year) + '-' + str(data_file.quarter))
+
+        # Paramiko need local file
+        paramiko_local_file = data_file.file.read()
+        with open(destination, 'wb') as f1:
+            f1.write(paramiko_local_file)
+            file_transfer_record.file_size = f1.tell()
+            file_transfer_record.file_shasum = hashlib.sha256(paramiko_local_file).hexdigest()
+            f1.close()
+
+        # Paramiko SSH connection requires private key as file
+        temp_key_file = write_key_to_file(local_key)
+        os.chmod(temp_key_file, 0o600)
+
+        # Create SFTP/SSH connection
+        transport = paramiko.SSHClient()
+        transport.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        pkey = paramiko.RSAKey.from_private_key_file(temp_key_file)
+        transport.connect(server_address,
+                          pkey=pkey,
+                          username=username,
+                          port=port,
+                          look_for_keys=False,
+                          disabled_algorithms={'pubkeys': ['rsa-sha2-512', 'rsa-sha2-256']})
+        # remove temp key file
+        os.remove(temp_key_file)
+        sftp = transport.open_sftp()
+
+        # Create remote directory
+        create_dir(settings.ACFTITAN_DIRECTORY, sftp_server=sftp)
+        create_dir(upper_directory_name, sftp_server=sftp)
+        create_dir(lower_directory_name, sftp_server=sftp)
+
+        # Put the file in SFTP server
+        sftp.put(destination, destination)
+
+        # Delete temp file
+        os.remove(destination)
+        logger.info('File {} has been successfully uploaded to {}'.format(destination, server_address))
+
+        # Add the log LegacyFileTransfer
+        file_transfer_record.result = LegacyFileTransfer.Result.COMPLETED
+        file_transfer_record.save()
+        transport.close()
+        return True
+
+    except Exception as e:
+        logger.error('Failed to upload {} with error:{}'.format(destination, e))
+        file_transfer_record.file_size = 0
+        file_transfer_record.result = LegacyFileTransfer.Result.ERROR
+        file_transfer_record.save()
+        transport.close()
+        return False