81 lines
2.0 KiB
Python
81 lines
2.0 KiB
Python
import os
|
|
import json
|
|
import threading
|
|
from enum import Enum
|
|
|
|
from colony.utils.files import atomic_write_json
|
|
|
|
STORAGE_DIR = "storage"
|
|
|
|
_write_lock = threading.Lock()
|
|
|
|
|
|
class Storage:
|
|
def __init__(self, storage_dir):
|
|
self.storage_dir = storage_dir
|
|
|
|
os.makedirs(self.storage_dir, exist_ok=True)
|
|
self._root = os.path.realpath(self.storage_dir)
|
|
|
|
def _resolve(self, name, allow_root=False):
|
|
if not isinstance(name, str) or "\0" in name:
|
|
raise ValueError("Invalid storage path")
|
|
|
|
path = os.path.realpath(os.path.join(self._root, name))
|
|
if os.path.commonpath([self._root, path]) != self._root:
|
|
raise ValueError("Storage path escapes the storage directory")
|
|
if path == self._root and not allow_root:
|
|
raise ValueError("Storage path must name a file")
|
|
return path
|
|
|
|
def get_path(self, filename):
|
|
return self._resolve(filename)
|
|
|
|
def ls(self, subdir):
|
|
path = self._resolve(subdir, allow_root=True)
|
|
|
|
if not os.path.isdir(path):
|
|
return []
|
|
|
|
return [name for name in os.listdir(path) if not name.startswith(".tmp-")]
|
|
|
|
def save(self, filename, data):
|
|
path = self.get_path(filename)
|
|
|
|
with _write_lock:
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
atomic_write_json(path, data, indent=2)
|
|
|
|
return path
|
|
|
|
def open(self, filename):
|
|
path = self.get_path(filename)
|
|
|
|
if not os.path.exists(path):
|
|
return None
|
|
|
|
with open(path, 'rb') as f:
|
|
return json.loads(f.read())
|
|
|
|
def exists(self, filename):
|
|
return os.path.exists(self.get_path(filename))
|
|
|
|
def delete(self, filename):
|
|
path = self.get_path(filename)
|
|
|
|
with _write_lock:
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
class DefaultStorage(Storage):
|
|
def __init__(self):
|
|
super().__init__(STORAGE_DIR)
|
|
|
|
|
|
class DEFAULT_STORAGE_FILES(Enum):
|
|
COLONY_SNAPSHOT = 'colony_snapshot.json'
|