# maubot - A plugin-based Matrix bot system.
# Copyright (C) 2019 Tulir Asokan
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from typing import Dict, List, Type, Tuple, Optional
from zipfile import ZipFile, BadZipFile
from time import time
import logging
import sys
import os
from ruamel.yaml import YAML, YAMLError
from packaging.version import Version
from mautrix.client.api.types.util import SerializerError
from ..lib.zipimport import zipimporter, ZipImportError
from ..plugin_base import Plugin
from ..config import Config
from .abc import PluginLoader, PluginClass, PluginMeta, IDConflictError
yaml = YAML()
class MaubotZipImportError(Exception):
pass
class MaubotZipMetaError(MaubotZipImportError):
pass
class MaubotZipPreLoadError(MaubotZipImportError):
pass
class MaubotZipLoadError(MaubotZipImportError):
pass
class ZippedPluginLoader(PluginLoader):
path_cache: Dict[str, 'ZippedPluginLoader'] = {}
log: logging.Logger = logging.getLogger("maubot.loader.zip")
trash_path: str = "delete"
directories: List[str] = []
path: str
meta: PluginMeta
main_class: str
main_module: str
_loaded: Type[PluginClass]
_importer: zipimporter
_file: ZipFile
def __init__(self, path: str) -> None:
super().__init__()
self.path = path
self.meta = None
self._loaded = None
self._importer = None
self._file = None
self._load_meta()
self._run_preload_checks(self._get_importer())
try:
existing = self.id_cache[self.meta.id]
raise IDConflictError(
f"Plugin with id {self.meta.id} already loaded from {existing.source}")
except KeyError:
pass
self.path_cache[self.path] = self
self.id_cache[self.meta.id] = self
self.log.debug(f"Preloaded plugin {self.meta.id} from {self.path}")
def to_dict(self) -> dict:
return {
**super().to_dict(),
"path": self.path
}
@classmethod
def get(cls, path: str) -> 'ZippedPluginLoader':
path = os.path.abspath(path)
try:
return cls.path_cache[path]
except KeyError:
return cls(path)
@property
def source(self) -> str:
return self.path
def __repr__(self) -> str:
return ("<ZippedPlugin "
f"path='{self.path}' "
f"meta={self.meta} "
f"loaded={self._loaded is not None}>")
async def read_file(self, path: str) -> bytes:
return self._file.read(path)
@staticmethod
def _read_meta(source) -> Tuple[ZipFile, PluginMeta]:
try:
file = ZipFile(source)
data = file.read("maubot.yaml")
except FileNotFoundError as e:
raise MaubotZipMetaError("Maubot plugin not found") from e
except BadZipFile as e:
raise MaubotZipMetaError("File is not a maubot plugin") from e
except KeyError as e:
raise MaubotZipMetaError("File does not contain a maubot plugin definition") from e
try:
meta_dict = yaml.load(data)
except (YAMLError, KeyError, IndexError, ValueError) as e:
raise MaubotZipMetaError("Maubot plugin definition file is not valid YAML") from e
try:
meta = PluginMeta.deserialize(meta_dict)
except SerializerError as e:
raise MaubotZipMetaError("Maubot plugin definition in file is invalid") from e
return file, meta
@classmethod
def verify_meta(cls, source) -> Tuple[str, Version]:
_, meta = cls._read_meta(source)
return meta.id, meta.version
def _load_meta(self) -> None:
file, meta = self._read_meta(self.path)
if self.meta and meta.id != self.meta.id:
raise MaubotZipMetaError("Maubot plugin ID changed during reload")
self.meta = meta
if "/" in meta.main_class:
self.main_module, self.main_class = meta.main_class.split("/")[:2]
else:
self.main_module = meta.modules[0]
self.main_class = meta.main_class
self._file = file
def _get_importer(self, reset_cache: bool = False) -> zipimporter:
try:
if not self._importer or self._importer.archive != self.path:
self._importer = zipimporter(self.path)
if reset_cache:
self._importer.reset_cache()
return self._importer
except ZipImportError as e:
raise MaubotZipMetaError("File not found or not a maubot plugin") from e
def _run_preload_checks(self, importer: zipimporter) -> None:
try:
code = importer.get_code(self.main_module.replace(".", "/"))
if self.main_class not in code.co_names:
raise MaubotZipPreLoadError(
f"Main class {self.main_class} not in {self.main_module}")
except ZipImportError as e:
raise MaubotZipPreLoadError(
f"Main module {self.main_module} not found in file") from e
for module in self.meta.modules:
try:
importer.find_module(module)
except ZipImportError as e:
raise MaubotZipPreLoadError(f"Module {module} not found in file") from e
async def load(self, reset_cache: bool = False) -> Type[PluginClass]:
try:
return self._load(reset_cache)
except MaubotZipImportError:
self.log.exception(f"Failed to load {self.meta.id} v{self.meta.version}")
raise
def _load(self, reset_cache: bool = False) -> Type[PluginClass]:
if self._loaded is not None and not reset_cache:
return self._loaded
self._load_meta()
importer = self._get_importer(reset_cache=reset_cache)
self._run_preload_checks(importer)
if reset_cache:
self.log.debug(f"Re-preloaded plugin {self.meta.id} from {self.path}")
for module in self.meta.modules:
try:
importer.load_module(module)
except ZipImportError:
raise MaubotZipLoadError(f"Module {module} not found in file")
except Exception:
raise MaubotZipLoadError(f"Failed to load module {module}")
try:
main_mod = sys.modules[self.main_module]
except KeyError as e:
raise MaubotZipLoadError(f"Main module {self.main_module} of plugin not found") from e
try:
plugin = getattr(main_mod, self.main_class)
except AttributeError as e:
raise MaubotZipLoadError(f"Main class {self.main_class} of plugin not found") from e
if not issubclass(plugin, Plugin):
raise MaubotZipLoadError("Main class of plugin does not extend maubot.Plugin")
self._loaded = plugin
self.log.debug(f"Loaded and imported plugin {self.meta.id} from {self.path}")
return plugin
async def reload(self, new_path: Optional[str] = None) -> Type[PluginClass]:
await self.unload()
if new_path is not None:
self.path = new_path
return await self.load(reset_cache=True)
async def unload(self) -> None:
for name, mod in list(sys.modules.items()):
if (getattr(mod, "__file__", "") or "").startswith(self.path):
del sys.modules[name]
self._loaded = None
self.log.debug(f"Unloaded plugin {self.meta.id} at {self.path}")
async def delete(self) -> None:
await self.unload()
try:
del self.path_cache[self.path]
except KeyError:
pass
try:
del self.id_cache[self.meta.id]
except KeyError:
pass
if self._importer:
self._importer.remove_cache()
self._importer = None
self._loaded = None
self.trash(self.path, reason="delete")
self.meta = None
self.path = None
@classmethod
def trash(cls, file_path: str, new_name: Optional[str] = None, reason: str = "error") -> None:
if cls.trash_path == "delete":
os.remove(file_path)
else:
new_name = new_name or f"{int(time())}-{reason}-{os.path.basename(file_path)}"
os.rename(file_path, os.path.abspath(os.path.join(cls.trash_path, new_name)))
@classmethod
def load_all(cls):
cls.log.debug("Preloading plugins...")
for directory in cls.directories:
for file in os.listdir(directory):
if not file.endswith(".mbp"):
continue
path = os.path.abspath(os.path.join(directory, file))
try:
cls.get(path)
except MaubotZipImportError:
cls.log.exception(f"Failed to load plugin at {path}, trashing...")
cls.trash(path)
except IDConflictError:
cls.log.error(f"Duplicate plugin ID at {path}, trashing...")
cls.trash(path)
def init(config: Config) -> None:
ZippedPluginLoader.trash_path = config["plugin_directories.trash"]
ZippedPluginLoader.directories = config["plugin_directories.load"]
ZippedPluginLoader.load_all()
