Commit initial changes to upload path

Author: TheByronHimes

config_schema.json

@@ -4,11 +4,18 @@
   "properties": {
     "max_concurrent_downloads": {
       "default": 5,
-      "description": "Number of parallel downloader tasks for file parts.",
+      "description": "Number of parallel download tasks for file parts.",
       "exclusiveMinimum": 0,
       "title": "Max Concurrent Downloads",
       "type": "integer"
     },
+    "max_concurrent_uploads": {
+      "default": 5,
+      "description": "Number of parallel upload tasks for file parts.",
+      "exclusiveMinimum": 0,
+      "title": "Max Concurrent Uploads",
+      "type": "integer"
+    },
     "max_retries": {
       "default": 5,
       "description": "Number of times to retry failed API calls.",
@@ -24,7 +31,7 @@
       "type": "integer"
     },
     "part_size": {
-      "default": 16777216,
+      "default": 67108864,
       "description": "The part size to use for download.",
       "exclusiveMinimum": 0,
       "title": "Part Size",

example_config.yaml

@@ -1,8 +1,9 @@
 exponential_backoff_max: 60
 max_concurrent_downloads: 5
+max_concurrent_uploads: 5
 max_retries: 2
 max_wait_time: 3600
-part_size: 16777216
+part_size: 67108864
 retry_status_codes:
 - 408
 - 500

src/ghga_connector/cli.py

@@ -22,18 +22,18 @@
 import typer
 
 from ghga_connector import exceptions
-from ghga_connector.config import CONFIG, set_runtime_config
 from ghga_connector.constants import C4GH
-from ghga_connector.core import CLIMessageDisplay, async_client
-from ghga_connector.core.main import async_download, decrypt_file, upload_file
+from ghga_connector.core import CLIMessageDisplay
+from ghga_connector.core.main import async_download, async_upload, decrypt_file
 from ghga_connector.core.utils import modify_for_debug, strtobool
 
 cli = typer.Typer(no_args_is_help=True)
 
 
+@cli.command(no_args_is_help=True)
 def upload(  # noqa: PLR0913
     *,
-    file_id: str = typer.Option(..., help="The id of the file to upload"),
+    file_alias: str = typer.Option(..., help="The alias of the file to upload"),
     file_path: Path = typer.Option(..., help="The path to the file to upload"),
     my_public_key_path: Path = typer.Option(
         "./key.pub",
@@ -58,7 +58,7 @@ def upload(  # noqa: PLR0913
     modify_for_debug(debug)
     asyncio.run(
         async_upload(
-            file_id=file_id,
+            file_alias=file_alias,
             file_path=file_path,
             my_public_key_path=my_public_key_path,
             my_private_key_path=my_private_key_path,
@@ -67,26 +67,6 @@ def upload(  # noqa: PLR0913
     )
 
 
-async def async_upload(
-    file_id: str,
-    file_path: Path,
-    my_public_key_path: Path,
-    my_private_key_path: Path,
-    passphrase: str | None = None,
-):
-    """Upload a file asynchronously"""
-    async with async_client() as client, set_runtime_config(client=client):
-        await upload_file(
-            client=client,
-            file_id=file_id,
-            file_path=file_path,
-            my_public_key_path=my_public_key_path,
-            my_private_key_path=my_private_key_path,
-            passphrase=passphrase,
-            part_size=CONFIG.part_size,
-        )
-
-
 if strtobool(os.getenv("UPLOAD_ENABLED") or "false"):
     cli.command(no_args_is_help=True)(upload)
 

src/ghga_connector/config.py

@@ -79,7 +79,10 @@ class Config(BaseSettings):
     """Global Config Parameters"""
 
     max_concurrent_downloads: PositiveInt = Field(
-        default=5, description="Number of parallel downloader tasks for file parts."
+        default=5, description="Number of parallel download tasks for file parts."
+    )
+    max_concurrent_uploads: PositiveInt = Field(
+        default=5, description="Number of parallel upload tasks for file parts."
     )
     max_retries: NonNegativeInt = Field(
         default=MAX_RETRIES, description="Number of times to retry failed API calls."

src/ghga_connector/constants.py

@@ -16,7 +16,7 @@
 
 """Constants used throughout the core."""
 
-DEFAULT_PART_SIZE = 16 * 1024 * 1024
+DEFAULT_PART_SIZE = 64 * (1024**2)  # 64 MiB
 TIMEOUT = 60.0
 TIMEOUT_LONG = 5 * TIMEOUT + 10
 MAX_PART_NUMBER = 10000

src/ghga_connector/core/crypt/__init__.py

@@ -15,7 +15,6 @@
 #
 """Subpackage containing all encryption/decryption related functionality"""
 
-from .abstract_bases import Decryptor, Encryptor  # noqa: F401
 from .checksums import Checksums  # noqa: F401
 from .decryption import Crypt4GHDecryptor  # noqa: F401
 from .encryption import Crypt4GHEncryptor  # noqa: F401

src/ghga_connector/core/crypt/abstract_bases.py

@@ -54,3 +54,10 @@ def update_encrypted(self, part: bytes):
         """Update encrypted part checksums"""
         self._encrypted_md5.append(hashlib.md5(part, usedforsecurity=False).hexdigest())
         self._encrypted_sha256.append(hashlib.sha256(part).hexdigest())
+
+    def encrypted_checksum_for_s3(self) -> str:
+        """Formulate the expected encrypted checksum str (etag) stored by S3."""
+        concatenated_md5s = b"".join(bytes.fromhex(md5) for md5 in self._encrypted_md5)
+        object_md5 = hashlib.md5(concatenated_md5s, usedforsecurity=False).hexdigest()
+        num_parts = len(self._encrypted_md5)
+        return object_md5 + f"-{num_parts}"

src/ghga_connector/core/crypt/checksums.py

@@ -20,10 +20,8 @@
 import crypt4gh.keys
 import crypt4gh.lib
 
-from .abstract_bases import Decryptor
 
-
-class Crypt4GHDecryptor(Decryptor):
+class Crypt4GHDecryptor:
     """Convenience class to deal with Crypt4GH decryption"""
 
     def __init__(self, decryption_key_path: Path, passphrase: str | None):

src/ghga_connector/core/crypt/decryption.py

@@ -13,46 +13,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-"""This module contains Crypt4GH based encryption functionality"""
+"""Functionality to encrypt files in chunks with Crypt4GH before upload."""
 
 import base64
 import os
 from collections.abc import Generator
 from io import BufferedReader
-from pathlib import Path
+from typing import Any
 
 import crypt4gh.header
-import crypt4gh.keys
 import crypt4gh.lib
 from nacl.bindings import crypto_aead_chacha20poly1305_ietf_encrypt
+from pydantic import SecretBytes
 
 from ghga_connector.config import get_ghga_pubkey
-from ghga_connector.core import get_segments, read_file_parts
+from ghga_connector.core.crypt.checksums import Checksums
+from ghga_connector.core.file_operations import get_segments, read_file_parts
 
-from .abstract_bases import Encryptor
-from .checksums import Checksums
 
-
-class Crypt4GHEncryptor(Encryptor):
+class Crypt4GHEncryptor:
     """Handles on the fly encryption and checksum calculation"""
 
-    def __init__(
-        self,
-        part_size: int,
-        private_key_path: Path,
-        passphrase: str | None,
-        checksums: Checksums = Checksums(),
-        file_secret: bytes | None = None,
-    ):
-        self._encrypted_file_size = 0
-        self._checksums = checksums
+    def __init__(self, part_size: int, my_private_key: SecretBytes):
         self._part_size = part_size
-        self._private_key_path = private_key_path
+        self._my_private_key = my_private_key
         self._server_public_key = base64.b64decode(get_ghga_pubkey())
-        self._passphrase = passphrase
-        if file_secret is None:
-            file_secret = os.urandom(32)
-        self._file_secret = file_secret
+        self._file_secret = os.urandom(32)
+        self.checksums = Checksums()  # Updated as encryption takes place
+        self._encrypted_file_size = 0  # Updated as encryption takes place
 
     def _encrypt(self, part: bytes):
         """Encrypt file part using secret"""
@@ -64,7 +52,7 @@ def _encrypt(self, part: bytes):
 
         return b"".join(encrypted_segments), incomplete_segment
 
-    def _encrypt_segment(self, segment: bytes):
+    def _encrypt_segment(self, segment: bytes) -> bytes:
         """Encrypt one single segment"""
         nonce = os.urandom(12)
         encrypted_data = crypto_aead_chacha20poly1305_ietf_encrypt(
@@ -77,38 +65,31 @@ def _create_envelope(self) -> bytes:
         Gather file encryption/decryption secret and assemble a crypt4gh envelope using the
         server's private and the client's public key
         """
-        if self._passphrase:
-            private_key = crypt4gh.keys.get_private_key(
-                filepath=self._private_key_path, callback=lambda: self._passphrase
-            )
-        else:
-            private_key = crypt4gh.keys.get_private_key(
-                filepath=self._private_key_path, callback=None
-            )
-
-        keys = [(0, private_key, self._server_public_key)]
+        keys = [(0, self._my_private_key.get_secret_value(), self._server_public_key)]
         header_content = crypt4gh.header.make_packet_data_enc(0, self._file_secret)
         header_packets = crypt4gh.header.encrypt(header_content, keys)
         header_bytes = crypt4gh.header.serialize(header_packets)
-
         return header_bytes
 
     def get_encrypted_size(self) -> int:
         """Get file size after encryption, excluding envelope"""
         return self._encrypted_file_size
 
-    def process_file(self, file: BufferedReader) -> Generator[bytes, None, None]:
-        """Encrypt file parts and prepare for upload."""
+    def process_file(
+        self, file: BufferedReader
+    ) -> Generator[tuple[int, bytes], Any, None]:
+        """Encrypt file parts for upload, yielding a tuple of the part number and content."""
         unprocessed_bytes = b""
         upload_buffer = self._create_envelope()
-        update_encrypted = self._checksums.update_encrypted
 
         # get envelope size to adjust checksum buffers and encrypted content size
         envelope_size = len(upload_buffer)
 
-        for file_part in read_file_parts(file=file, part_size=self._part_size):
+        for part_number, file_part in enumerate(
+            read_file_parts(file=file, part_size=self._part_size), start=1
+        ):
             # process unencrypted
-            self._checksums.update_unencrypted(file_part)
+            self.checksums.update_unencrypted(file_part)
             unprocessed_bytes += file_part
 
             # encrypt in chunks
@@ -118,12 +99,12 @@ def process_file(self, file: BufferedReader) -> Generator[bytes, None, None]:
             # update checksums and yield if part size
             if len(upload_buffer) >= self._part_size:
                 current_part = upload_buffer[: self._part_size]
-                if self._checksums.encrypted_is_empty():
-                    update_encrypted(current_part[envelope_size:])
+                if self.checksums.encrypted_is_empty():
+                    self.checksums.update_encrypted(current_part[envelope_size:])
                 else:
-                    update_encrypted(current_part)
+                    self.checksums.update_encrypted(current_part)
                 self._encrypted_file_size += self._part_size
-                yield current_part
+                yield part_number, current_part
                 upload_buffer = upload_buffer[self._part_size :]
 
         # process dangling bytes
@@ -132,14 +113,16 @@ def process_file(self, file: BufferedReader) -> Generator[bytes, None, None]:
 
         while len(upload_buffer) >= self._part_size:
             current_part = upload_buffer[: self._part_size]
-            update_encrypted(current_part)
+            self.checksums.update_encrypted(current_part)
             self._encrypted_file_size += self._part_size
-            yield current_part
+            part_number += 1
+            yield part_number, current_part
             upload_buffer = upload_buffer[self._part_size :]
 
         if upload_buffer:
-            update_encrypted(upload_buffer)
+            self.checksums.update_encrypted(upload_buffer)
             self._encrypted_file_size += len(upload_buffer)
-            yield upload_buffer
+            part_number += 1
+            yield part_number, upload_buffer
 
         self._encrypted_file_size -= envelope_size