Add evaluations

scripts/pentest_fixtures.py

@@ -1,8 +1,12 @@
+import hashlib
 import json
+import os
 import random
 import secrets
 import string
 import subprocess
+import tarfile
+import time
 from contextlib import contextmanager
 from io import BytesIO
 from pathlib import Path
@@ -16,7 +20,6 @@
 from django.contrib.auth import get_user_model
 from django.core.files.base import ContentFile
 from django.core.files.uploadedfile import InMemoryUploadedFile
-from django.db import transaction
 from faker import Faker
 from PIL import Image as PILImage
 from PIL import ImageDraw, ImageFont
@@ -30,7 +33,7 @@
     ComponentInterfaceValue,
 )
 from grandchallenge.direct_messages.models import Conversation, DirectMessage
-from grandchallenge.evaluation.models import Method, Phase
+from grandchallenge.evaluation.models import Method, Phase, Submission
 from grandchallenge.evaluation.utils import SubmissionKindChoices
 from grandchallenge.invoices.models import Invoice
 from grandchallenge.pages.models import Page
@@ -44,8 +47,9 @@
 from grandchallenge.verifications.models import Verification
 from grandchallenge.workstations.models import Workstation
 
+NUM_ARCHIVE_ITEMS = 1
+
 
-@transaction.atomic
 def run():
     with _demo_container() as container:
         create_challenges(container=container)
@@ -132,20 +136,33 @@ def create_challenge(*, challenge_num, container):
         payment_status=Invoice.PaymentStatusChoices.PAID,
     )
 
+    participant_1_algorithm = create_algorithm(
+        container=container,
+        title=f"Challenge {challenge_num}, Participant 1, Algorithm",
+        editor=participant_1,
+    )
+    participant_2_algorithm = create_algorithm(
+        container=container,
+        title=f"Challenge {challenge_num}, Participant 2, Algorithm",
+        editor=participant_2,
+    )
+
     for ii, give_algorithm_editors_job_view_permissions in enumerate(
         (True, False)
     ):
         phase_num = ii + 1
 
+        phase_title = f"Challenge {challenge_num}, Phase {phase_num}"
+
         archive = _create_phase_archive(
             creator=admin,
             interfaces=_get_inputs(),
-            title=f"Challenge {challenge_num}, Phase {phase_num} Test Set",
+            title=f"{phase_title}, Test Set",
         )
 
         p = Phase.objects.create(
             challenge=c,
-            title=f"Challenge {challenge_num} Phase {phase_num}",
+            title=phase_title,
             algorithm_time_limit=300,
             submission_kind=SubmissionKindChoices.ALGORITHM,
             archive=archive,
@@ -162,17 +179,40 @@ def create_challenge(*, challenge_num, container):
         with open(container, "rb") as f:
             m.image.save("method_io.tar", ContentFile(f.read()))
 
-    create_algorithm(
-        container=container,
-        title=f"Challenge {challenge_num}, Participant 1, Algorithm",
-        editor=participant_1,
-    )
-    create_algorithm(
-        container=container,
-        title=f"Challenge {challenge_num}, Participant 2, Algorithm",
-        editor=participant_2,
+        await_method_and_image_validation()
+
+        Submission.objects.create(
+            creator=participant_1,
+            phase=p,
+            algorithm_image=participant_1_algorithm.active_image,
+            comment=f"{phase_title}, Participant 1, Submission",
+        )
+        Submission.objects.create(
+            creator=participant_2,
+            phase=p,
+            algorithm_image=participant_2_algorithm.active_image,
+            comment=f"{phase_title}, Participant 2, Submission",
+        )
+
+
+def await_method_and_image_validation():
+    print(
+        ">>> Waiting for Method and AlgorithmImage validation to complete..."
     )
 
+    while True:
+        if (
+            AlgorithmImage.objects.exclude(
+                import_status=AlgorithmImage.ImportStatusChoices.COMPLETED
+            ).exists()
+            or Method.objects.exclude(
+                import_status=Method.ImportStatusChoices.COMPLETED
+            ).exists()
+        ):
+            time.sleep(1)
+        else:
+            break
+
 
 def _get_inputs():
     return ComponentInterface.objects.filter(
@@ -208,7 +248,7 @@ def create_archive(*, archive_num):
     add_archive_item(archive=archive, interfaces=_get_inputs())
 
 
-def _create_phase_archive(*, creator, interfaces, title, items=5):
+def _create_phase_archive(*, creator, interfaces, title):
     archive = Archive.objects.create(
         title=title,
         logo=create_image_with_text(text=f"{title}, Logo"),
@@ -218,7 +258,7 @@ def _create_phase_archive(*, creator, interfaces, title, items=5):
     )
     archive.add_editor(creator)
 
-    for _ in range(items):
+    for _ in range(NUM_ARCHIVE_ITEMS):
         add_archive_item(archive=archive, interfaces=interfaces)
 
     return archive
@@ -341,8 +381,15 @@ def create_algorithm(
 
     algorithm_image = AlgorithmImage(creator=editor, algorithm=algorithm)
 
-    with open(container, "rb") as f:
-        algorithm_image.image.save("algorithm_io.tar", ContentFile(f.read()))
+    with add_env_var_to_container(
+        tarball_path=container, new_env_var=f"TITLE={title!r}"
+    ) as image:
+        with open(image, "rb") as f:
+            algorithm_image.image.save(
+                "algorithm_io.tar", ContentFile(f.read())
+            )
+
+    return algorithm
 
 
 @contextmanager
@@ -519,3 +566,86 @@ def create_direct_messages(users):
             )
             if unread:
                 message.unread_by.set({*combination} - {sender})
+
+
+def extract_tarball(tarball_path, extract_path):
+    with tarfile.open(tarball_path, "r") as tar:
+        tar.extractall(path=extract_path)
+
+
+def calculate_sha256(file_path):
+    sha256 = hashlib.sha256()
+    with open(file_path, "rb") as file:
+        while chunk := file.read(8192):
+            sha256.update(chunk)
+    return sha256.hexdigest()
+
+
+def modify_config_json(config_json_path, new_env_var):
+    with open(config_json_path) as file:
+        config = json.load(file)
+
+    env_var_key, env_var_value = new_env_var.split("=", 1)
+    new_env_entry = f"{env_var_key}={env_var_value}"
+
+    if "Env" in config["config"]:
+        config["config"]["Env"].append(new_env_entry)
+    else:
+        config["config"]["Env"] = [new_env_entry]
+
+    with open(config_json_path, "w") as file:
+        json.dump(config, file)
+
+    return config
+
+
+def update_manifest(manifest_path, config_digest, config_size):
+    with open(manifest_path) as file:
+        manifest = json.load(file)
+
+    manifest[0]["Config"] = f"{config_digest}.json"
+    manifest[0]["ConfigSize"] = config_size
+
+    with open(manifest_path, "w") as file:
+        json.dump(manifest, file)
+
+
+def rename_file_with_digest(original_path, new_digest):
+    new_path = os.path.join(
+        os.path.dirname(original_path), f"{new_digest}.json"
+    )
+    os.rename(original_path, new_path)
+    return new_path
+
+
+def create_tarball(source_dir, tarball_path):
+    with tarfile.open(tarball_path, "w") as tar:
+        for root, _, files in os.walk(source_dir):
+            for file in files:
+                fullpath = os.path.join(root, file)
+                arcname = os.path.relpath(fullpath, start=source_dir)
+                tar.add(fullpath, arcname=arcname)
+
+
+@contextmanager
+def add_env_var_to_container(tarball_path, new_env_var):
+    with TemporaryDirectory() as work_dir:
+        extract_tarball(tarball_path, work_dir)
+
+        manifest_path = os.path.join(work_dir, "manifest.json")
+        with open(manifest_path) as file:
+            manifest = json.load(file)
+
+        config_json_path = os.path.join(work_dir, manifest[0]["Config"])
+        modify_config_json(config_json_path, new_env_var)
+
+        config_digest = calculate_sha256(config_json_path)
+        config_size = os.path.getsize(config_json_path)
+
+        rename_file_with_digest(config_json_path, config_digest)
+        update_manifest(manifest_path, config_digest, config_size)
+
+        with NamedTemporaryFile() as new_tarball_path:
+            create_tarball(work_dir, new_tarball_path)
+
+            yield new_tarball_path