refactoring

3 files changed, 115 insertions, 102 deletions

diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index 50f737a6..7347a733 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -57,8 +57,7 @@ from .. import BaseMsd
 from .. import MsdFileReader
 from .. import MsdFileWriter
 
-from . import fs
-
+from .storage import Storage
 from .drive import Drive
 
 
@@ -67,9 +66,15 @@ from .drive import Drive
 class _DriveImage:
     name: str
     path: str
-    size: int
     complete: bool
     in_storage: bool
+    size: int = dataclasses.field(default=0)
+
+    def __post_init__(self) -> None:
+        try:
+            object.__setattr__(self, "size", max(os.path.getsize(self.path), 0))
+        except Exception as err:
+            get_logger().warning("Can't get size of file %s: %s", self.path, err)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -151,16 +156,13 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         self.__write_chunk_size = write_chunk_size
         self.__sync_chunk_size = sync_chunk_size
 
-        self.__storage_path = os.path.normpath(storage_path)
-        self.__images_path = os.path.join(self.__storage_path, "images")
-        self.__meta_path = os.path.join(self.__storage_path, "meta")
-
         self.__remount_cmd = remount_cmd
 
         self.__initial_image: str = initial["image"]
         self.__initial_cdrom: bool = initial["cdrom"]
 
         self.__drive = Drive(gadget, instance=0, lun=0)
+        self.__storage = Storage(storage_path)
 
         self.__reader: (MsdFileReader | None) = None
         self.__writer: (MsdFileWriter | None) = None
@@ -206,7 +208,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 if self.__writer:
                     # При загрузке файла показываем актуальную статистику вручную
                     storage["uploading"] = self.__writer.get_state()
-                    space = fs.get_fs_space(self.__storage_path, fatal=False)
+                    space = self.__storage.get_space(fatal=False)
                     if space:
                         storage.update(dataclasses.asdict(space))
                 else:
@@ -338,7 +340,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         if self.__state.vd.connected or self.__drive.get_image_path():
                             raise MsdConnectedError()
 
-                        path = os.path.join(self.__images_path, name)
+                        path = self.__storage.get_image_path(name)
                         if name not in self.__state.storage.images or not os.path.exists(path):
                             raise MsdUnknownImageError()
 
@@ -369,12 +371,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                         if self.__state.vd.connected or self.__drive.get_image_path():
                             raise MsdConnectedError()
 
-                        path = os.path.join(self.__images_path, name)
+                        path = self.__storage.get_image_path(name)
                         if name in self.__state.storage.images or os.path.exists(path):
                             raise MsdImageExistsError()
 
                         await self.__remount_rw(True)
-                        self.__set_image_complete(name, False)
+                        self.__storage.set_image_complete(name, False)
 
                         self.__writer = await MsdFileWriter(
                             notifier=self.__notifier,
@@ -386,7 +388,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     self.__notifier.notify()
                     yield self.__writer
-                    self.__set_image_complete(name, self.__writer.is_complete())
+                    self.__storage.set_image_complete(name, self.__writer.is_complete())
 
                 finally:
                     if remove_incomplete and self.__writer and not self.__writer.is_complete():
@@ -424,7 +426,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
             await self.__remount_rw(True)
             os.remove(image.path)
-            self.__set_image_complete(name, False)
+            self.__storage.set_image_complete(name, False)
             await self.__remount_rw(False, fatal=False)
 
     # =====
@@ -453,9 +455,10 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                     await asyncio.sleep(5)
 
                 with Inotify() as inotify:
-                    inotify.watch(self.__images_path, InotifyMask.ALL_MODIFY_EVENTS)
-                    inotify.watch(self.__meta_path, InotifyMask.ALL_MODIFY_EVENTS)
-                    for path in self.__drive.get_watchable_paths():
+                    for path in [
+                        *self.__storage.get_watchable_paths(),
+                        *self.__drive.get_watchable_paths(),
+                    ]:
                         inotify.watch(path, InotifyMask.ALL_MODIFY_EVENTS)
 
                     # После установки вотчеров еще раз проверяем стейт, чтобы ничего не потерять
@@ -522,7 +525,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
     async def __setup_initial(self) -> None:
         if self.__initial_image:
             logger = get_logger(0)
-            path = os.path.join(self.__images_path, self.__initial_image)
+            path = self.__storage.get_image_path(self.__initial_image)
             if os.path.exists(path):
                 logger.info("Setting up initial image %r ...", self.__initial_image)
                 try:
@@ -538,19 +541,14 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     def __get_storage_state(self) -> _StorageState:
         images: dict[str, _DriveImage] = {}
-        for name in os.listdir(self.__images_path):
-            path = os.path.join(self.__images_path, name)
-            if os.path.exists(path):
-                size = fs.get_file_size(path)
-                if size >= 0:
-                    images[name] = _DriveImage(
-                        name=name,
-                        path=path,
-                        size=size,
-                        complete=self.__is_image_complete(name),
-                        in_storage=True,
-                    )
-        space = fs.get_fs_space(self.__storage_path, fatal=True)
+        for name in self.__storage.get_images():
+            images[name] = _DriveImage(
+                name=name,
+                path=self.__storage.get_image_path(name),
+                complete=self.__storage.is_image_complete(name),
+                in_storage=True,
+            )
+        space = self.__storage.get_space(fatal=True)
         assert space
         return _StorageState(
             size=space.size,
@@ -563,12 +561,11 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
         path = self.__drive.get_image_path()
         if path:
             name = os.path.basename(path)
-            in_storage = (os.path.dirname(path) == self.__images_path)
+            in_storage = self.__storage.is_image_path_in_storage(path)
             image = _DriveImage(
                 name=name,
                 path=path,
-                size=max(fs.get_file_size(path), 0),
-                complete=(self.__is_image_complete(name) if in_storage else True),
+                complete=(self.__storage.is_image_complete(name) if in_storage else True),
                 in_storage=in_storage,
             )
         return _DriveState(
@@ -579,19 +576,6 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
     # =====
 
-    def __is_image_complete(self, name: str) -> bool:
-        return os.path.exists(os.path.join(self.__meta_path, name + ".complete"))
-
-    def __set_image_complete(self, name: str, flag: bool) -> None:
-        path = os.path.join(self.__meta_path, name + ".complete")
-        if flag:
-            open(path, "w").close()  # pylint: disable=consider-using-with
-        else:
-            if os.path.exists(path):
-                os.remove(path)
-
-    # =====
-
     async def __remount_rw(self, rw: bool, fatal: bool=True) -> None:
         if not (await aiohelpers.remount("MSD", self.__remount_cmd, rw)):
             if fatal:
diff --git a/kvmd/plugins/msd/otg/fs.py b/kvmd/plugins/msd/otg/fs.py
deleted file mode 100644
index d5018190..00000000
--- a/kvmd/plugins/msd/otg/fs.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# ========================================================================== #
-#                                                                            #
-#    KVMD - The main PiKVM daemon.                                           #
-#                                                                            #
-#    Copyright (C) 2018-2022  Maxim Devaev <[email protected]>               #
-#                                                                            #
-#    This program is free software: you can redistribute it and/or modify    #
-#    it under the terms of the GNU 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 General Public License for more details.                            #
-#                                                                            #
-#    You should have received a copy of the GNU General Public License       #
-#    along with this program.  If not, see <https://www.gnu.org/licenses/>.  #
-#                                                                            #
-# ========================================================================== #
-
-
-import os
-import dataclasses
-
-from ....logging import get_logger
-
-
-# =====
[email protected](frozen=True)
-class FsSpace:
-    size: int
-    free: int
-
-
-# =====
-def get_file_size(path: str) -> int:
-    try:
-        return os.path.getsize(path)
-    except Exception as err:
-        get_logger().warning("Can't get size of file %s: %s", path, err)
-        return -1
-
-
-def get_fs_space(path: str, fatal: bool) -> (FsSpace | None):
-    try:
-        st = os.statvfs(path)
-    except Exception as err:
-        if fatal:
-            raise
-        get_logger().warning("Can't get free space of filesystem %s: %s", path, err)
-        return None
-    return FsSpace(
-        size=(st.f_blocks * st.f_frsize),
-        free=(st.f_bavail * st.f_frsize),
-    )
diff --git a/kvmd/plugins/msd/otg/storage.py b/kvmd/plugins/msd/otg/storage.py
new file mode 100644
index 00000000..28ab0fa4
--- /dev/null
+++ b/kvmd/plugins/msd/otg/storage.py
@@ -0,0 +1,85 @@
+# ========================================================================== #
+#                                                                            #
+#    KVMD - The main PiKVM daemon.                                           #
+#                                                                            #
+#    Copyright (C) 2018-2022  Maxim Devaev <[email protected]>               #
+#                                                                            #
+#    This program is free software: you can redistribute it and/or modify    #
+#    it under the terms of the GNU 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 General Public License for more details.                            #
+#                                                                            #
+#    You should have received a copy of the GNU General Public License       #
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.  #
+#                                                                            #
+# ========================================================================== #
+
+
+import os
+import dataclasses
+
+from ....logging import get_logger
+
+
+# =====
[email protected](frozen=True)
+class StorageSpace:
+    size: int
+    free: int
+
+
+class Storage:
+    def __init__(self, path: str) -> None:
+        self.__path = path
+        self.__images_path = os.path.join(self.__path, "images")
+        self.__meta_path = os.path.join(self.__path, "meta")
+
+    def get_watchable_paths(self) -> list[str]:
+        return [self.__images_path, self.__meta_path]
+
+    def get_images(self) -> list[str]:
+        images: list[str] = []
+        for name in os.listdir(self.__images_path):
+            path = os.path.join(self.__images_path, name)
+            if os.path.exists(path):
+                try:
+                    if os.path.getsize(path) >= 0:
+                        images.append(name)
+                except Exception:
+                    pass
+        return images
+
+    def get_image_path(self, name: str) -> str:
+        return os.path.join(self.__images_path, name)
+
+    def is_image_path_in_storage(self, path: str) -> bool:
+        return (os.path.dirname(path) == self.__images_path)
+
+    def is_image_complete(self, name: str) -> bool:
+        return os.path.exists(os.path.join(self.__meta_path, name + ".complete"))
+
+    def set_image_complete(self, name: str, flag: bool) -> None:
+        path = os.path.join(self.__meta_path, name + ".complete")
+        if flag:
+            open(path, "w").close()  # pylint: disable=consider-using-with
+        else:
+            if os.path.exists(path):
+                os.remove(path)
+
+    def get_space(self, fatal: bool) -> (StorageSpace | None):
+        try:
+            st = os.statvfs(self.__path)
+        except Exception as err:
+            if fatal:
+                raise
+            get_logger().warning("Can't get free space of filesystem %s: %s", self.__path, err)
+            return None
+        return StorageSpace(
+            size=(st.f_blocks * st.f_frsize),
+            free=(st.f_bavail * st.f_frsize),
+        )