Files
Auto-GPT/autogpts/swifty/forge/sdk/workspace.py
2023-09-15 21:27:02 +02:00

77 lines
2.4 KiB
Python

import abc
import os
import typing
from pathlib import Path
class Workspace(abc.ABC):
@abc.abstractclassmethod
def __init__(self, base_path: str) -> None:
self.base_path = base_path
@abc.abstractclassmethod
def read(self, task_id: str, path: str) -> bytes:
pass
@abc.abstractclassmethod
def write(self, task_id: str, path: str, data: bytes) -> None:
pass
@abc.abstractclassmethod
def delete(
self, task_id: str, path: str, directory: bool = False, recursive: bool = False
) -> None:
pass
@abc.abstractclassmethod
def exists(self, task_id: str, path: str) -> bool:
pass
@abc.abstractclassmethod
def list(self, task_id: str, path: str) -> typing.List[str]:
pass
class LocalWorkspace(Workspace):
def __init__(self, base_path: str):
self.base_path = Path(base_path).resolve()
def _resolve_path(self, task_id: str, path: str) -> Path:
abs_path = (self.base_path / task_id / path).resolve()
if not str(abs_path).startswith(str(self.base_path)):
print("Error")
raise ValueError("Directory traversal is not allowed!")
abs_path.parent.mkdir(parents=True, exist_ok=True)
return abs_path
def read(self, task_id: str, path: str) -> bytes:
with open(self._resolve_path(task_id, path), "rb") as f:
return f.read()
def write(self, task_id: str, path: str, data: bytes) -> None:
file_path = self._resolve_path(task_id, path)
with open(file_path, "wb") as f:
f.write(data)
def delete(
self, task_id: str, path: str, directory: bool = False, recursive: bool = False
) -> None:
path = self.base_path / task_id / path
resolved_path = self._resolve_path(task_id, path)
if directory:
if recursive:
os.rmdir(resolved_path)
else:
os.removedirs(resolved_path)
else:
os.remove(resolved_path)
def exists(self, task_id: str, path: str) -> bool:
path = self.base_path / task_id / path
return self._resolve_path(task_id, path).exists()
def list(self, task_id: str, path: str) -> typing.List[str]:
path = self.base_path / task_id / path
base = self._resolve_path(task_id, path)
return [str(p.relative_to(self.base_path / task_id)) for p in base.iterdir()]