102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
import hashlib
|
|
import io
|
|
import os
|
|
import json
|
|
import tarfile
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from colony.utils.files import atomic_write_json
|
|
|
|
STORAGES_DIR = "storages"
|
|
|
|
_write_lock = threading.Lock()
|
|
|
|
|
|
class Storage:
|
|
SHARED = False
|
|
|
|
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
|
|
|
|
def dir_checksum(self):
|
|
h = hashlib.sha256()
|
|
path = Path(self.storage_dir)
|
|
|
|
for file in sorted(p for p in path.rglob("*") if p.is_file() and not p.name.startswith(".tmp-")):
|
|
h.update(str(file.relative_to(path)).encode())
|
|
with file.open("rb") as f:
|
|
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
|
h.update(chunk)
|
|
|
|
return h.hexdigest()
|
|
|
|
def archive(self):
|
|
buffer = io.BytesIO()
|
|
|
|
with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
|
|
for file in sorted(p for p in Path(self.storage_dir).rglob("*") if p.is_file() and not p.name.startswith(".tmp-")):
|
|
tar.add(file, arcname=file.relative_to(self.storage_dir).as_posix())
|
|
|
|
return buffer.getvalue()
|
|
|
|
def restore(self, data):
|
|
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
|
|
tar.extractall(self.storage_dir, filter="data")
|