Adds password support for upload task (#137)

* Adds password support for upload task

* Adds usage of password

* Renames password to key

* Renames password to key

* Renames password to key

* Non working version

* Makes secret key owrking

* Changes encoding steps

* Updates submodule

* Embedding providers (#136)

* forward platform information

* Adds agreement, gdpr comliant flsgs and embedding provider support

* Adds created by to embeddings

* Merges embedding logic for different embedding types and add support more embedding providers

* Change order of embedding platforms

* Added platform as part of the Encoder type

* Adds modularization for recreation of embeddings

* Refactors recreation of embeddings to solve dependency problems

* Embedding recreation for adding new records

* Fixes embedding deletion

* Fixes small things

* Adds notification for organization change

* Adds better state management for reupload of records

* Adds optimized state management for embeddings

* Standardize call to recreate embeddings

* Resolves PR comments

* Removes print statement

* Adds commit for embedding creation

* Adds handling for missing tokenization

* Adds logic to infer embedding information out of old projects

* Adds order of project transfer so that source code can be replaced by new embedding name

* Changes call logic of agreements

* Resolves first few PR comments

* Resolves PR comments

* Adds new term text

* Changed terms text slightly

* Adds link and placeholder

* Added link to embedding type

* Update controller/transfer/project_transfer_manager.py

Co-authored-by: JWittmeyer <[email protected]>

* Resolves typo

* Resolves typo

* Submodules merge

* Drone

---------

Co-authored-by: Johannes Hötter <[email protected]>
Co-authored-by: JWittmeyer <[email protected]>
Co-authored-by: Lina <[email protected]>

* Merges and adds key support for upload

* Fixes key handling

* Standardizes behavior how to export

* Fixes export of projects

* Removes print

* Removes key of upload task after upload

* Adds clean up logic for keys in upload tasks

* Adds cleanup of files in tmp

* Cleans file after project export

* Clean up code

* Embedding providers (#136)

* forward platform information

* Adds agreement, gdpr comliant flsgs and embedding provider support

* Adds created by to embeddings

* Merges embedding logic for different embedding types and add support more embedding providers

* Change order of embedding platforms

* Added platform as part of the Encoder type

* Adds modularization for recreation of embeddings

* Refactors recreation of embeddings to solve dependency problems

* Embedding recreation for adding new records

* Fixes embedding deletion

* Fixes small things

* Adds notification for organization change

* Adds better state management for reupload of records

* Adds optimized state management for embeddings

* Standardize call to recreate embeddings

* Resolves PR comments

* Removes print statement

* Adds commit for embedding creation

* Adds handling for missing tokenization

* Adds logic to infer embedding information out of old projects

* Adds order of project transfer so that source code can be replaced by new embedding name

* Changes call logic of agreements

* Resolves first few PR comments

* Resolves PR comments

* Adds new term text

* Changed terms text slightly

* Adds link and placeholder

* Added link to embedding type

* Update controller/transfer/project_transfer_manager.py

Co-authored-by: JWittmeyer <[email protected]>

* Resolves typo

* Resolves typo

* Submodules merge

* Drone

---------

Co-authored-by: Johannes Hötter <[email protected]>
Co-authored-by: JWittmeyer <[email protected]>
Co-authored-by: Lina <[email protected]>

* Update start

* Update start

* Adds check for secret key on startup

* Sets warning on default secret key

* Adds handling for bad password in file import

* Rewrites additional apt packages

* Removes print

* Removes unnessecary global usage and formats code

* Resolves PR comments

* Resolves PR comments

* Resolves PR comments

* Submodules merge

---------

Co-authored-by: Johannes Hötter <[email protected]>
Co-authored-by: JWittmeyer <[email protected]>
Co-authored-by: Lina <[email protected]>

Author: 4 people

Dockerfile

@@ -2,6 +2,9 @@ FROM kernai/refinery-parent-images:v1.11.0-common
 
 WORKDIR /app
 
+# used for encryption and zipping of files
+RUN apt-get update && apt-get install -y libc6-dev zlib1g gcc --no-install-recommends
+
 COPY requirements.txt .
 
 RUN pip3 install --no-cache-dir -r requirements.txt

alembic/versions/217dbe11c5d1_adds_key_for_upload.py

@@ -0,0 +1,28 @@
+"""Adds key for upload
+
+Revision ID: 217dbe11c5d1
+Revises: 1a25c862801f
+Create Date: 2023-06-21 14:55:16.523327
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '217dbe11c5d1'
+down_revision = '1a25c862801f'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('upload_task', sa.Column('key', sa.LargeBinary(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('upload_task', 'key')
+    # ### end Alembic commands ###

api/transfer.py

@@ -1,16 +1,15 @@
 import logging
 import traceback
 import time
-from typing import Any, List
+from typing import Optional
 
 from controller import organization
-from controller.embedding import util as embedding_util
-from controller.embedding import connector as embedding_connector
 from starlette.endpoints import HTTPEndpoint
 from starlette.responses import PlainTextResponse, JSONResponse
 from controller.embedding.manager import recreate_embeddings
 
 from controller.transfer.labelstudio import import_preperator
+from exceptions.exceptions import BadPasswordError
 from submodules.s3 import controller as s3
 from submodules.model.business_objects import (
     attribute,
@@ -69,6 +68,14 @@ async def post(self, request) -> PlainTextResponse:
         is_global_update = True if task.file_type == "project" else False
         try:
             init_file_import(task, project_id, is_global_update)
+        except BadPasswordError:
+            file_import_error_handling(
+                task,
+                project_id,
+                is_global_update,
+                enums.NotificationType.BAD_PASSWORD_DURING_IMPORT,
+                print_traceback=False,
+            )
         except Exception:
             file_import_error_handling(task, project_id, is_global_update)
         notification.send_organization_update(
@@ -268,13 +275,19 @@ def init_file_import(task: UploadTask, project_id: str, is_global_update: bool)
 
 
 def file_import_error_handling(
-    task: UploadTask, project_id: str, is_global_update: bool
+    task: UploadTask,
+    project_id: str,
+    is_global_update: bool,
+    notification_type: Optional[NotificationType] = None,
+    print_traceback: bool = True,
 ) -> None:
     general.rollback()
     task.state = enums.UploadStates.ERROR.value
     general.commit()
+    if not notification_type:
+        notification_type = NotificationType.IMPORT_FAILED
     create_notification(
-        NotificationType.IMPORT_FAILED,
+        notification_type,
         task.user_id,
         task.project_id,
         task.file_type,
@@ -284,13 +297,17 @@ def file_import_error_handling(
             task,
         )
     )
-    print(traceback.format_exc(), flush=True)
+    if print_traceback:
+        print(traceback.format_exc(), flush=True)
+
     notification.send_organization_update(
         project_id, f"file_upload:{str(task.id)}:state:{task.state}", is_global_update
     )
 
 
-def __recalculate_missing_attributes_and_embeddings(project_id: str, user_id: str) -> None:
+def __recalculate_missing_attributes_and_embeddings(
+    project_id: str, user_id: str
+) -> None:
     __calculate_missing_attributes(project_id, user_id)
     recreate_embeddings(project_id)
 
@@ -384,4 +401,3 @@ def __calculate_missing_attributes(project_id: str, user_id: str) -> None:
             message="calculate_attribute:finished:all",
         )
         general.remove_and_refresh_session(ctx_token, False)
-        

app.py

@@ -21,6 +21,7 @@
 from graphql_api import schema
 from controller.task_queue.task_queue import init_task_queue
 from controller.project.manager import check_in_deletion_projects
+from util import security, clean_up
 
 
 logging.basicConfig(level=logging.DEBUG)
@@ -57,3 +58,6 @@
 
 init_task_queue()
 check_in_deletion_projects()
+security.check_secret_key()
+clean_up.clean_up_database()
+clean_up.clean_up_disk()

controller/notification/notification_data.py

@@ -121,6 +121,13 @@
         "page": enums.Pages.SETTINGS.value,
         "docs": enums.DOCS.UPLOADING_DATA.value,
     },
+    enums.NotificationType.BAD_PASSWORD_DURING_IMPORT.value: {
+        "message_template": "Bad password for zip file",
+        "title": "Data import",
+        "level": enums.Notification.ERROR.value,
+        "page": enums.Pages.OVERVIEW.value,
+        "docs": enums.DOCS.CREATING_PROJECTS.value,
+    },
     enums.NotificationType.INFORMATION_SOURCE_STARTED.value: {
         "message_template": "Started heuristic @@arg@@.",
         "title": "Heuristic execution",

controller/transfer/manager.py

@@ -3,7 +3,6 @@
 import json
 import traceback
 from typing import Any, List, Optional, Dict
-import zipfile
 
 from controller.transfer import export_parser
 from controller.transfer.knowledge_base_transfer_manager import (
@@ -36,14 +35,15 @@
 from submodules.s3 import controller as s3
 import pandas as pd
 from datetime import datetime
-from util import notification
+from util import notification, security, file
 from sqlalchemy.sql import text as sql_text
 from controller.labeling_task import manager as labeling_task_manager
 from controller.labeling_task_label import manager as labeling_task_label_manager
 from submodules.model.business_objects import record_label_association as rla
 from controller.task_queue import manager as task_queue_manager
 from submodules.model.enums import TaskType, RecordTokenizationScope
 
+
 from util.notification import create_notification
 
 logging.basicConfig(level=logging.DEBUG)
@@ -57,9 +57,17 @@ def get_upload_credentials_and_id(
     file_type: str,
     file_import_options: str,
     upload_type: str,
+    key: Optional[str] = None,
 ):
+    key = security.encrypt(key)
     task = upload_task_manager.create_upload_task(
-        str(user_id), project_id, file_name, file_type, file_import_options, upload_type
+        str(user_id),
+        project_id,
+        file_name,
+        file_type,
+        file_import_options,
+        upload_type,
+        key,
     )
     org_id = organization.get_id_by_project_id(project_id)
     return s3.get_upload_credentials_and_id(org_id, project_id + "/" + str(task.id))
@@ -174,7 +182,10 @@ def export_records(
 
 
 def prepare_record_export(
-    project_id: str, user_id: str, export_options: Optional[Dict[str, Any]] = None
+    project_id: str,
+    user_id: str,
+    export_options: Optional[Dict[str, Any]] = None,
+    key: Optional[str] = None,
 ) -> None:
     records_by_options_query_data = get_records_by_options_query_data(
         project_id, export_options
@@ -187,7 +198,7 @@ def prepare_record_export(
     file_path, file_name = export_parser.parse(
         project_id, final_query, mapping_dict, extraction_appends, export_options
     )
-    zip_path, file_name = __write_file_to_zip(file_path)
+    zip_path, file_name = file.file_to_zip(file_path, key)
     org_id = organization.get_id_by_project_id(project_id)
     prefixed_path = f"{project_id}/download/{user_id}/record_export_"
     file_name_download = prefixed_path + file_name
@@ -233,7 +244,10 @@ def export_knowledge_base(project_id: str, base_id: str) -> str:
 
 
 def prepare_project_export(
-    project_id: str, user_id: str, export_options: Dict[str, bool]
+    project_id: str,
+    user_id: str,
+    export_options: Dict[str, bool],
+    key: Optional[str] = None,
 ) -> bool:
     org_id = organization.get_id_by_project_id(project_id)
     objects = s3.get_bucket_objects(org_id, project_id + "/download/project_export_")
@@ -242,36 +256,19 @@ def prepare_project_export(
 
     data = get_project_export_dump(project_id, user_id, export_options)
     file_name_base = "project_export_" + datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
-    file_name_local = file_name_base + ".zip"
-    file_name_download = project_id + "/download/" + file_name_local
-    __write_json_data_to_zip(data, file_name_base)
-    s3.upload_object(org_id, file_name_download, file_name_local)
+    json_file_path = file.text_to_json_file(data, file_name_base)
+    zip_path, zip_name = file.file_to_zip(json_file_path, key)
+    file_name_download = f"{project_id}/download/{zip_name}"
+    s3.upload_object(org_id, file_name_download, zip_path)
     notification.send_organization_update(project_id, "project_export")
 
-    if os.path.exists(file_name_local):
-        os.remove(file_name_local)
+    if os.path.exists(json_file_path):
+        os.remove(json_file_path)
+    if os.path.exists(zip_path):
+        os.remove(zip_path)
     return True
 
 
-def __write_file_to_zip(file_path: str) -> str:
-    base_name = os.path.basename(file_path)
-    file_name = base_name + ".zip"
-    zip_path = f"{file_path}.zip"
-    zipfile.ZipFile(zip_path, mode="w").write(file_path, base_name)
-    return zip_path, file_name
-
-
-def __write_json_data_to_zip(dumped_json: str, base_file_name: str) -> None:
-    with zipfile.ZipFile(
-        base_file_name + ".zip",
-        mode="w",
-        compression=zipfile.ZIP_DEFLATED,
-        compresslevel=9,
-    ) as zip_file:
-        zip_file.writestr(base_file_name + ".json", data=dumped_json)
-        zip_file.testzip()
-
-
 def last_project_export_credentials(project_id: str) -> str:
     return __get_last_export_credentials(project_id, "/download/project_export_")