Fix pathing issues

This commit is contained in:
SwiftyOS
2023-09-28 12:29:03 +02:00
parent ece0f9189f
commit 4f15b1c582
2 changed files with 22 additions and 5 deletions

View File

@@ -37,11 +37,15 @@ class LocalWorkspace(Workspace):
self.base_path = Path(base_path).resolve() self.base_path = Path(base_path).resolve()
def _resolve_path(self, task_id: str, path: str) -> Path: def _resolve_path(self, task_id: str, path: str) -> Path:
path = path if not path.startswith("/") else path[1:]
abs_path = (self.base_path / task_id / path).resolve() abs_path = (self.base_path / task_id / path).resolve()
if not str(abs_path).startswith(str(self.base_path)): if not str(abs_path).startswith(str(self.base_path)):
print("Error") print("Error")
raise ValueError(f"Directory traversal is not allowed! - {abs_path}") raise ValueError(f"Directory traversal is not allowed! - {abs_path}")
try:
abs_path.parent.mkdir(parents=True, exist_ok=True) abs_path.parent.mkdir(parents=True, exist_ok=True)
except FileExistsError:
pass
return abs_path return abs_path
def read(self, task_id: str, path: str) -> bytes: def read(self, task_id: str, path: str) -> bytes:

View File

@@ -1,5 +1,7 @@
import json import json
import logging
import os import os
import pathlib
import time import time
from typing import Any, Dict, Optional from typing import Any, Dict, Optional
@@ -14,6 +16,8 @@ from agbenchmark.agent_protocol_client import (
from agbenchmark.agent_protocol_client.models.step import Step from agbenchmark.agent_protocol_client.models.step import Step
from agbenchmark.utils.data_types import ChallengeData from agbenchmark.utils.data_types import ChallengeData
LOG = logging.getLogger(__name__)
async def run_api_agent( async def run_api_agent(
task: ChallengeData, config: Dict[str, Any], artifacts_location: str, timeout: int task: ChallengeData, config: Dict[str, Any], artifacts_location: str, timeout: int
@@ -63,13 +67,22 @@ async def copy_agent_artifacts_into_temp_folder(api_instance, task_id):
artifacts = await api_instance.list_agent_task_artifacts(task_id=task_id) artifacts = await api_instance.list_agent_task_artifacts(task_id=task_id)
for artifact in artifacts.artifacts: for artifact in artifacts.artifacts:
# current absolute path of the directory of the file # current absolute path of the directory of the file
directory_location = TEMP_FOLDER_ABS_PATH directory_location = pathlib.Path(TEMP_FOLDER_ABS_PATH)
if artifact.relative_path: if artifact.relative_path:
directory_location = os.path.dirname(directory_location / artifact.relative_path) path = (
artifact.relative_path
if not artifact.relative_path.startswith("/")
else artifact.relative_path[1:]
)
directory_location = pathlib.Path(
os.path.dirname(directory_location / path)
)
LOG.info(f"Creating directory {directory_location}")
os.mkdir(directory_location, recursive=True, exist_ok=True) directory_location.mkdir(parents=True, exist_ok=True)
file_path = directory_location / artifact.file_name file_path = directory_location / artifact.file_name
LOG.info(f"Writing file {file_path}")
with open(file_path, "wb") as f: with open(file_path, "wb") as f:
content = await api_instance.download_agent_task_artifact( content = await api_instance.download_agent_task_artifact(
task_id=task_id, artifact_id=artifact.artifact_id task_id=task_id, artifact_id=artifact.artifact_id