From 2dd11e9c596c264a3d00c5b85bd8e378c6b49da3 Mon Sep 17 00:00:00 2001
From: Eric Zhu <ekzhu@users.noreply.github.com>
Date: Sat, 13 Jan 2024 10:42:44 -0800
Subject: [PATCH] [Core] Sanitize filename before using it as docker image tag.
 Fix #1069 (#1127)

* fix #1069

* uncomment skips

* use importlib to check package

* better code

---------

Co-authored-by: Chi Wang <wang.chi@microsoft.com>
---
 autogen/code_utils.py | 28 ++++++++++++++++++++++++++--
 test/test_code.py     | 40 ++++++++++++++++++++++++++++++++--------
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/autogen/code_utils.py b/autogen/code_utils.py
index 299a39fba9d0..f468c3d1bfdc 100644
--- a/autogen/code_utils.py
+++ b/autogen/code_utils.py
@@ -2,6 +2,7 @@
 import os
 import pathlib
 import re
+import string
 import subprocess
 import sys
 import time
@@ -224,6 +225,29 @@ def _cmd(lang):
     raise NotImplementedError(f"{lang} not recognized in code execution")
 
 
+def _sanitize_filename_for_docker_tag(filename: str) -> str:
+    """Convert a filename to a valid docker tag.
+    See https://docs.docker.com/engine/reference/commandline/tag/ for valid tag
+    format.
+
+    Args:
+        filename (str): The filename to be converted.
+
+    Returns:
+        str: The sanitized Docker tag.
+    """
+    # Replace any character not allowed with an underscore
+    allowed_chars = set(string.ascii_letters + string.digits + "_.-")
+    sanitized = "".join(char if char in allowed_chars else "_" for char in filename)
+
+    # Ensure it does not start with a period or a dash
+    if sanitized.startswith(".") or sanitized.startswith("-"):
+        sanitized = "_" + sanitized[1:]
+
+    # Truncate if longer than 128 characters
+    return sanitized[:128]
+
+
 def execute_code(
     code: Optional[str] = None,
     timeout: Optional[int] = None,
@@ -379,7 +403,7 @@ def execute_code(
     cmd = [
         "sh",
         "-c",
-        f"{_cmd(lang)} {filename}; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str}",
+        f'{_cmd(lang)} "{filename}"; exit_code=$?; echo -n {exit_code_str}; echo -n $exit_code; echo {exit_code_str}',
     ]
     # create a docker container
     container = client.containers.run(
@@ -403,7 +427,7 @@ def execute_code(
     # get the container logs
     logs = container.logs().decode("utf-8").rstrip()
     # commit the image
-    tag = filename.replace("/", "")
+    tag = _sanitize_filename_for_docker_tag(filename)
     container.commit(repository="python", tag=tag)
     # remove the container
     container.remove()
diff --git a/test/test_code.py b/test/test_code.py
index ec8ada2c9062..e7306125de2a 100644
--- a/test/test_code.py
+++ b/test/test_code.py
@@ -1,3 +1,4 @@
+import importlib.metadata
 import os
 import sys
 import unittest
@@ -22,6 +23,21 @@
 here = os.path.abspath(os.path.dirname(__file__))
 
 
+def is_package_installed(package_name):
+    """Check if a package is installed. This is a preferred way to check if a
+    package is installed or not than doing a try-catch around an import
+    because it avoids name conflict with local modules and
+    code execution in the imported module."""
+    try:
+        importlib.metadata.version(package_name)
+        return True
+    except importlib.metadata.PackageNotFoundError:
+        return False
+
+
+docker_package_installed = is_package_installed("docker")
+
+
 # def test_find_code():
 #     try:
 #         import openai
@@ -293,10 +309,6 @@ def scrape(url):
     assert len(codeblocks) == 1 and codeblocks[0] == ("", "source setup.sh")
 
 
-@pytest.mark.skipif(
-    sys.platform in ["darwin"],
-    reason="do not run on MacOS",
-)
 def test_execute_code(use_docker=None):
     try:
         import docker
@@ -338,15 +350,27 @@ def test_execute_code(use_docker=None):
     assert isinstance(image, str) or docker is None or os.path.exists("/.dockerenv") or use_docker is False
 
 
+@pytest.mark.skipif(docker_package_installed is False, reason="docker package not installed")
+def test_execute_code_with_custom_filename_on_docker():
+    exit_code, msg, image = execute_code("print('hello world')", filename="tmp/codetest.py", use_docker=True)
+    assert exit_code == 0 and msg == "hello world\n", msg
+    assert image == "python:tmp_codetest.py"
+
+
+@pytest.mark.skipif(docker_package_installed is False, reason="docker package not installed")
+def test_execute_code_with_misformed_filename_on_docker():
+    exit_code, msg, image = execute_code(
+        "print('hello world')", filename="tmp/codetest.py (some extra information)", use_docker=True
+    )
+    assert exit_code == 0 and msg == "hello world\n", msg
+    assert image == "python:tmp_codetest.py__some_extra_information_"
+
+
 def test_execute_code_raises_when_code_and_filename_are_both_none():
     with pytest.raises(AssertionError):
         execute_code(code=None, filename=None)
 
 
-@pytest.mark.skipif(
-    sys.platform in ["darwin"],
-    reason="do not run on MacOS",
-)
 def test_execute_code_nodocker():
     test_execute_code(use_docker=False)