415 lines
14 KiB
Python
415 lines
14 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
# SPDX-FileCopyrightText: 2023 KDE e.V.
|
|
# SPDX-FileContributor: Ingo Klöcker <dev@ingo-kloecker.de>
|
|
|
|
import dataclasses
|
|
import json
|
|
import logging
|
|
import os
|
|
import stat
|
|
import time
|
|
from pathlib import Path, PurePath, PurePosixPath
|
|
from tempfile import TemporaryDirectory
|
|
from typing import List, Optional
|
|
|
|
import sftpnotary.log
|
|
from sftpnotary import authentication, util
|
|
from sftpnotary.exceptions import Error
|
|
from sftpnotary.sftp import ConnectionLostError, SFTPClient
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MetaData:
|
|
def asDict(self) -> dict:
|
|
return dataclasses.asdict(self)
|
|
|
|
def asJson(self) -> str:
|
|
# disable escaping of non-ASCII characters
|
|
return json.dumps(dataclasses.asdict(self), ensure_ascii=False)
|
|
|
|
@classmethod
|
|
def fromDict(cls, d: dict):
|
|
fields = dataclasses.fields(cls)
|
|
initKwargs = {field.name: d[field.name] for field in fields if field.name in d}
|
|
unknownFields = sorted(list(set(d.keys()) - set(initKwargs.keys())))
|
|
if unknownFields:
|
|
log.debug(f"Ignored the following unknown entries in the dict: {', '.join(unknownFields)}")
|
|
return cls(**initKwargs)
|
|
|
|
@classmethod
|
|
def fromJson(cls, s: str):
|
|
return cls.fromDict(json.loads(s))
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Request(MetaData):
|
|
token: Optional[str] = None
|
|
projectPath: Optional[str] = None
|
|
branch: Optional[str] = None
|
|
|
|
def validate(self):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SingleFileRequest(Request):
|
|
fileName: Optional[str] = None
|
|
|
|
def validate(self):
|
|
util.validateFileName(self.fileName)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Response(MetaData):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SignFileResponse(Response):
|
|
signedFileName: Optional[str] = None
|
|
|
|
|
|
class Task:
|
|
def __init__(self, id: str, basePath: PurePath):
|
|
self.id = id
|
|
self.basePath = basePath
|
|
self.queue = None
|
|
|
|
def path(self) -> PurePath:
|
|
return self.basePath / self.id
|
|
|
|
def requestPath(self) -> PurePath:
|
|
return self.path() / "request.json"
|
|
|
|
def responsePath(self) -> PurePath:
|
|
return self.path() / "response.json"
|
|
|
|
def putRequest(self, request):
|
|
self.queue.uploadJson(request.asJson(), self.requestPath())
|
|
|
|
def fetchRequest(self, requestClass):
|
|
data = self.queue.fetchJson(self.requestPath())
|
|
|
|
log.debug(f"Request: {data}")
|
|
return requestClass.fromJson(data)
|
|
|
|
def putResponse(self, response):
|
|
self.queue.uploadJson(response.asJson(), self.responsePath())
|
|
|
|
def fetchResponse(self, responseClass):
|
|
data = self.queue.fetchJson(self.responsePath())
|
|
|
|
log.debug(f"Response: {data}")
|
|
return responseClass.fromJson(data)
|
|
|
|
|
|
def _ensureFolderExists(sftp: SFTPClient, path: PurePath):
|
|
try:
|
|
sftp.stat(str(path))
|
|
return
|
|
except FileNotFoundError:
|
|
log.debug("Folder '%s' does not exist. Creating ...", path)
|
|
try:
|
|
sftp.mkdir(str(path))
|
|
except OSError as e:
|
|
raise Error(f"Failed to create folder '{path}' ({e!r})")
|
|
|
|
|
|
class TaskQueue:
|
|
def __init__(self, sftp: SFTPClient):
|
|
self.sftp = sftp
|
|
self.newPath = PurePosixPath("tmp")
|
|
self.pendingPath = PurePosixPath("queue")
|
|
self.donePath = PurePosixPath("done")
|
|
self.trashPath = PurePosixPath("trash")
|
|
_ensureFolderExists(self.sftp, self.newPath)
|
|
_ensureFolderExists(self.sftp, self.pendingPath)
|
|
_ensureFolderExists(self.sftp, self.donePath)
|
|
_ensureFolderExists(self.sftp, self.trashPath)
|
|
|
|
def createTask(self):
|
|
timestamp = time.strftime("%Y%m%dT%H%M%S", time.gmtime())
|
|
if "CI_JOB_ID" in os.environ:
|
|
clientId = f"{os.environ['CI_PROJECT_PATH_SLUG']}_{os.environ['CI_JOB_ID']}"
|
|
else:
|
|
clientId = os.uname().nodename
|
|
task = Task(id=f"{timestamp}-{clientId}", basePath=self.newPath)
|
|
task.queue = self
|
|
self.sftp.mkdir(str(task.path()))
|
|
return task
|
|
|
|
def commitTask(self, task: Task):
|
|
log.info(f"Committing task '{task.id}' ...")
|
|
if task.basePath != self.newPath:
|
|
raise Error(f"Task '{task.id}' not in folder for new tasks (task.path: {task.path()})")
|
|
|
|
try:
|
|
targetPath = self.pendingPath / task.id
|
|
self.sftp.posix_rename(str(task.path()), str(targetPath))
|
|
task.basePath = self.pendingPath
|
|
except OSError as e:
|
|
raise Error(f"Failed to commit task '{task.id}' ({e!r})")
|
|
return task
|
|
|
|
def _getFirstTaskInQueue(self):
|
|
entries = sorted(self.sftp.listdir(str(self.pendingPath)))
|
|
if len(entries) > 0:
|
|
task = Task(entries[0], self.pendingPath)
|
|
task.queue = self
|
|
return task
|
|
|
|
def getNextTask(self):
|
|
while True:
|
|
log.debug("Checking queue for next task ...")
|
|
task = self._getFirstTaskInQueue()
|
|
if task is not None:
|
|
log.debug(f"Found task {task.id}")
|
|
yield task
|
|
time.sleep(5)
|
|
|
|
def takeTask(self, task: Task, targetPath: PurePath):
|
|
log.info(f"Taking task {task.id} ...")
|
|
try:
|
|
self.sftp.posix_rename(str(task.path()), str(targetPath / task.id))
|
|
task.basePath = targetPath
|
|
return task
|
|
except FileNotFoundError:
|
|
# likely no error, but another worker took the task
|
|
log.info(f"Failed to take task '{task.id}'.")
|
|
return
|
|
|
|
def uploadJson(self, data: str, remotePath: PurePath):
|
|
self.sftp.uploadBytes(data.encode(), remotePath)
|
|
|
|
def fetchJson(self, remotePath: PurePath) -> str:
|
|
return self.sftp.downloadBytes(remotePath).decode()
|
|
|
|
def waitForCompletion(self, task: Task):
|
|
while True:
|
|
log.debug(f"Checking if task '{task.id}' is done ...")
|
|
try:
|
|
attrs = self.sftp.stat(str(self.donePath / task.id)) # raises FileNotFoundError
|
|
log.debug(f"Got attributes: {attrs}")
|
|
task.basePath = self.donePath
|
|
return task
|
|
except FileNotFoundError:
|
|
pass
|
|
except ConnectionLostError:
|
|
log.warning("Lost connection to SFTP server. Reconnecting ...")
|
|
self.sftp.reconnect()
|
|
|
|
time.sleep(5)
|
|
|
|
def setDone(self, task: Task):
|
|
log.info(f"Setting task '{task.id}' to done ...")
|
|
|
|
try:
|
|
self.sftp.posix_rename(str(task.path()), str(self.donePath / task.id))
|
|
task.basePath = self.donePath
|
|
return task
|
|
except OSError as e:
|
|
raise Error(f"Failed to set task '{task.id}' to done ({e!r})")
|
|
|
|
def removeTask(self, task: Task):
|
|
log.info(f"Moving task '{task.id}' to trash ...")
|
|
try:
|
|
self.sftp.posix_rename(str(task.path()), str(self.trashPath / task.id))
|
|
task.basePath = self.trashPath
|
|
return task
|
|
except OSError as e:
|
|
raise Error(f"Failed to move task '{task.id}' to trash ({e!r})")
|
|
|
|
def cleanUp(self, taskFolder: PurePath, maxAgeInDays: int = 0):
|
|
log.debug(f"Cleaning up task folder '{taskFolder}' ...")
|
|
now = time.time()
|
|
cutOffTime = now - maxAgeInDays * 24 * 60 * 60
|
|
entriesAttributes = self.sftp.listdir_attr(taskFolder)
|
|
for entryAttributes in entriesAttributes:
|
|
entryName = entryAttributes.filename
|
|
if not stat.S_ISDIR(entryAttributes.st_mode):
|
|
log.debug(f"Removing unexpected file '{entryName}'")
|
|
self.sftp.remove(taskFolder / entryName)
|
|
elif entryAttributes.st_mtime < cutOffTime:
|
|
log.debug(f"Removing task '{entryName}' ...")
|
|
self.sftp.recursiveRemoveDirectory(taskFolder / entryName)
|
|
|
|
def cleanUpTaskFolders(self, maxAgeInDays: int = 1):
|
|
# remove old tasks from all of the standard task folders
|
|
self.cleanUp(self.newPath, maxAgeInDays)
|
|
self.cleanUp(self.pendingPath, maxAgeInDays)
|
|
self.cleanUp(self.donePath, maxAgeInDays)
|
|
self.cleanUp(self.trashPath, maxAgeInDays)
|
|
|
|
|
|
def ProgressLogger(message: str):
|
|
def printProgress(current: int, total: int):
|
|
log.debug(message, current, total)
|
|
|
|
return printProgress
|
|
|
|
|
|
class Job:
|
|
def __init__(self, sftp: SFTPClient):
|
|
self.sftp = sftp
|
|
self._taskQueue = None
|
|
self._task = None
|
|
self.success: bool = False
|
|
self.request: Optional[Request] = None # created by the actual job
|
|
self.logFiles: List[Path] = []
|
|
|
|
@property
|
|
def taskQueue(self):
|
|
if self._taskQueue is None:
|
|
self._taskQueue = TaskQueue(self.sftp)
|
|
return self._taskQueue
|
|
|
|
@property
|
|
def task(self):
|
|
if self._task is None:
|
|
self._task = self.taskQueue.createTask()
|
|
return self._task
|
|
|
|
def listLogFiles(self) -> List[PurePath]:
|
|
logFiles = []
|
|
for attr in self.sftp.listdir_iter(str(self.task.path())):
|
|
if attr.filename.endswith(".log"):
|
|
logFiles.append(self.task.path() / attr.filename)
|
|
return logFiles
|
|
|
|
def downloadLogFiles(self, destinationFolder: Path) -> None:
|
|
self.logFiles.clear()
|
|
for logFile in self.listLogFiles():
|
|
localPath = destinationFolder / logFile.name
|
|
self.sftp.download(logFile, localPath)
|
|
self.logFiles.append(localPath)
|
|
|
|
|
|
class SignSingleFileJob(Job):
|
|
def __init__(self, sftp: SFTPClient, filePath: Path):
|
|
super().__init__(sftp)
|
|
|
|
self.localFilePath = Path(filePath)
|
|
|
|
def uploadUnsignedFile(self):
|
|
localPath = self.localFilePath
|
|
remotePath = self.task.path() / self.request.fileName
|
|
self.sftp.upload(localPath, remotePath)
|
|
|
|
|
|
class Worker:
|
|
def __init__(self, workerId: str, sftp: SFTPClient):
|
|
self.workerId = workerId
|
|
self.sftp = sftp
|
|
self._processingPath = None
|
|
self._taskQueue = None
|
|
|
|
@property
|
|
def processingPath(self):
|
|
if self._processingPath is None:
|
|
path = PurePosixPath(f"processing-{self.workerId}")
|
|
_ensureFolderExists(self.sftp, path)
|
|
self._processingPath = path
|
|
return self._processingPath
|
|
|
|
@property
|
|
def taskQueue(self):
|
|
if self._taskQueue is None:
|
|
self._taskQueue = TaskQueue(self.sftp)
|
|
return self._taskQueue
|
|
|
|
def start(self):
|
|
# remove leftovers in the worker's processing folder
|
|
self.taskQueue.cleanUp(self.processingPath)
|
|
# remove old tasks from the different task folders
|
|
self.taskQueue.cleanUpTaskFolders()
|
|
|
|
while True:
|
|
try:
|
|
self._runTaskLoop()
|
|
except ConnectionLostError:
|
|
log.warning("Lost connection to SFTP server. Reconnecting ...")
|
|
self.sftp.reconnect()
|
|
|
|
def _runTaskLoop(self):
|
|
for task in self.taskQueue.getNextTask():
|
|
task = self.taskQueue.takeTask(task, self.processingPath)
|
|
if task is not None:
|
|
processor = self.processorClass(self.sftp)
|
|
processor.process(task)
|
|
self.taskQueue.setDone(task)
|
|
# remove old tasks from the different task folders
|
|
self.taskQueue.cleanUpTaskFolders()
|
|
|
|
|
|
class TaskProcessor:
|
|
tries = 1
|
|
|
|
def __init__(self, sftp: SFTPClient):
|
|
self.sftp = sftp
|
|
self.task: Optional[Task] = None
|
|
self.request: Optional[Request] = None
|
|
|
|
def process(self, task: Task):
|
|
self.task = task
|
|
with TemporaryDirectory() as tmpdirname:
|
|
debugFilePath = Path(tmpdirname) / "task-debug.log"
|
|
logFilePath = Path(tmpdirname) / "task.log"
|
|
with sftpnotary.log.captureTaskDebugLog(debugFilePath), sftpnotary.log.captureTaskLog(logFilePath):
|
|
self.processInner()
|
|
self.sftp.upload(debugFilePath, self.task.path() / debugFilePath.name)
|
|
self.sftp.upload(logFilePath, self.task.path() / logFilePath.name)
|
|
|
|
def processInner(self):
|
|
log.debug(f"Fetching request of task '{self.task.id}' ...")
|
|
try:
|
|
self.request = self.task.fetchRequest(self.requestClass)
|
|
except Exception as e:
|
|
log.error(f"Error: Fetching request of task '{self.task.id}' failed.")
|
|
log.debug("Error details:", exc_info=e)
|
|
return
|
|
|
|
try:
|
|
self.request.validate()
|
|
except Exception as e:
|
|
log.error(f"Error: Request of task '{self.task.id}' is invalid.")
|
|
log.debug("Error details:", exc_info=e)
|
|
return
|
|
|
|
if authentication.authentication is not None:
|
|
log.debug("Authenticating request ...")
|
|
try:
|
|
authentication.authentication.authenticate(request=self.request)
|
|
except Exception as e:
|
|
log.error(f"Error: Authenticating request of task '{self.task.id}' failed.")
|
|
log.debug("Error details:", exc_info=e)
|
|
return
|
|
|
|
log.info(f"Processing task '{self.task.id}' ...")
|
|
|
|
done = False
|
|
tries = self.tries
|
|
while not done and tries > 0:
|
|
try:
|
|
response = self.doProcess()
|
|
if response is None:
|
|
response = Response()
|
|
self.task.putResponse(response)
|
|
done = True
|
|
except Exception as e:
|
|
tries -= 1
|
|
if tries > 0:
|
|
log.error("Error: Processing task '%s' failed with %r. Retrying ...", self.task.id, e)
|
|
else:
|
|
log.error("Error: Processing task '%s' failed with %r. Giving up.", self.task.id, e)
|
|
log.debug("Error details:", exc_info=e)
|
|
if not done:
|
|
log.debug("Marking task '%s' as failed.", self.task.id)
|
|
# mark as failed
|
|
pass
|
|
|
|
log.info(f"Finished processing task '{self.task.id}' ...")
|
|
|
|
def doProcess(self):
|
|
raise NotImplementedError("You need to use a suitable subclass!")
|