import importlib
|
import requests
|
import os
|
from os import path
|
|
from app import models, pending_actions
|
from app.plugins.views import TaskView
|
from app.plugins.worker import run_function_async
|
from app.plugins import get_current_plugin
|
|
from worker.celery import app
|
from rest_framework.response import Response
|
from rest_framework import status
|
|
from .platform_helper import get_all_platforms, get_platform_by_name
|
|
class ImportFolderTaskView(TaskView):
|
def post(self, request, project_pk=None, pk=None):
|
task = self.get_and_check_task(request, pk)
|
|
# Read form data
|
folder_url = request.data.get('selectedFolderUrl', None)
|
platform_name = request.data.get('platform', None)
|
|
# Make sure both values are set
|
if folder_url == None or platform_name == None:
|
return Response({'error': 'Folder URL and platform name must be set.'}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Fetch the platform by name
|
platform = get_platform_by_name(platform_name)
|
|
# Make sure that the platform actually exists
|
if platform == None:
|
return Response({'error': 'Failed to find a platform with the name \'{}\''.format(platform_name)}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Verify that the folder url is valid
|
if platform.verify_folder_url(folder_url) == None:
|
return Response({'error': 'Invalid URL'}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Get the files from the folder
|
files = platform.import_from_folder(folder_url)
|
|
# Update the task with the new information
|
task.console += "Importing {} images...\n".format(len(files))
|
task.images_count = len(files)
|
task.pending_action = pending_actions.IMPORT
|
task.save()
|
|
# Associate the folder url with the project and task
|
combined_id = "{}_{}".format(project_pk, pk)
|
get_current_plugin().get_global_data_store().set_string(combined_id, folder_url)
|
|
# Start importing the files in the background
|
serialized = [file.serialize() for file in files]
|
run_function_async(import_files, task.id, serialized)
|
|
return Response({}, status=status.HTTP_200_OK)
|
|
class CheckUrlTaskView(TaskView):
|
def get(self, request, project_pk=None, pk=None):
|
|
# Assert that task exists
|
self.get_and_check_task(request, pk)
|
|
# Check if there is an imported url associated with the project and task
|
combined_id = "{}_{}".format(project_pk, pk)
|
folder_url = get_current_plugin().get_global_data_store().get_string(combined_id, default = None)
|
|
if folder_url == None:
|
return Response({}, status=status.HTTP_200_OK)
|
else:
|
return Response({'folder_url': folder_url}, status=status.HTTP_200_OK)
|
|
class PlatformsVerifyTaskView(TaskView):
|
def get(self, request, platform_name):
|
# Read the form data
|
folder_url = request.GET.get('folderUrl', None)
|
|
# Fetch the platform by name
|
platform = get_platform_by_name(platform_name)
|
|
# Make sure that the platform actually exists
|
if platform == None:
|
return Response({'error': 'Failed to find a platform with the name \'{}\''.format(platform_name)}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Verify that the folder url is valid
|
folder = platform.verify_folder_url(folder_url)
|
if folder == None:
|
return Response({'error': 'Invalid URL'}, status=status.HTTP_400_BAD_REQUEST)
|
|
# Return the folder
|
return Response({'folder': folder.serialize()}, status=status.HTTP_200_OK)
|
|
|
class PlatformsTaskView(TaskView):
|
def get(self, request):
|
# Fetch and return all platforms
|
platforms = get_all_platforms()
|
return Response({'platforms': [platform.serialize(user = request.user) for platform in platforms]}, status=status.HTTP_200_OK)
|
|
|
def import_files(task_id, files):
|
import requests
|
from app import models
|
from app.plugins import logger
|
from app.security import path_traversal_check
|
|
def download_file(task, file):
|
path = path_traversal_check(task.task_path(file['name']), task.task_path())
|
download_stream = requests.get(file['url'], stream=True, timeout=60)
|
|
with open(path, 'wb') as fd:
|
for chunk in download_stream.iter_content(4096):
|
fd.write(chunk)
|
|
logger.info("Will import {} files".format(len(files)))
|
task = models.Task.objects.get(pk=task_id)
|
task.create_task_directories()
|
task.save()
|
|
try:
|
downloaded_total = 0
|
for file in files:
|
download_file(task, file)
|
task.check_if_canceled()
|
models.Task.objects.filter(pk=task.id).update(upload_progress=(float(downloaded_total) / float(len(files))))
|
downloaded_total += 1
|
|
except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
|
raise NodeServerError(e)
|
|
task.refresh_from_db()
|
task.pending_action = None
|
task.processing_time = 0
|
task.partial = False
|
task.images_count = len(task.scan_images())
|
task.save()
|
