From 3adf4755749e97e7c834aa51bc90f0e34e1b6eec Mon Sep 17 00:00:00 2001 From: themanyfaceddemon Date: Mon, 2 Sep 2024 17:40:31 +0300 Subject: [PATCH] Done --- DMBotNetwork/__init__.py | 4 +- DMBotNetwork/side/side_client.py | 105 +++++++++++++++++++++++------- DMBotNetwork/side/side_server.py | 17 ++--- DMBotNetwork/units/client_unit.py | 14 ++-- DMBotNetwork/utils/server_db.py | 4 +- Tests/ServerDB.py | 50 ++++++++++++++ Tests/__init__.py | 1 - Tests/bruh.py | 1 - setup.py | 4 +- 9 files changed, 154 insertions(+), 46 deletions(-) create mode 100644 Tests/ServerDB.py delete mode 100644 Tests/bruh.py diff --git a/DMBotNetwork/__init__.py b/DMBotNetwork/__init__.py index c313d33..b746004 100644 --- a/DMBotNetwork/__init__.py +++ b/DMBotNetwork/__init__.py @@ -1,5 +1,5 @@ -from side import Client, Server -from utils import ServerDB +from .side import Client, Server +from .utils import ServerDB __all__ = ["Client", "Server", "ServerDB"] __version__ = "0.1.0" diff --git a/DMBotNetwork/side/side_client.py b/DMBotNetwork/side/side_client.py index 909f6b4..54564d7 100644 --- a/DMBotNetwork/side/side_client.py +++ b/DMBotNetwork/side/side_client.py @@ -2,10 +2,12 @@ import inspect import logging from asyncio import StreamReader, StreamWriter +from pathlib import Path from typing import Any, Dict, Optional, get_type_hints import msgpack -from utils import NetCode + +from ..utils import NCAnyType, NCLogType, NetCode logger = logging.getLogger("DMBotNetwork Client") @@ -14,9 +16,11 @@ class Client: _network_methods: Dict[str, Any] = {} _ear_task: Optional[asyncio.Task] = None # lol + _viva_alp: bool = True _login: str = None _password: str = None - _is_downloaded: bool = True + _content_path: Optional[Path] = None + _server_name: Optional[str] = None _is_connected: bool = False _is_auth: bool = False @@ -77,6 +81,14 @@ async def connect(cls, host, port) -> None: def is_connected(cls) -> bool: return cls._is_auth and cls._is_connected + @classmethod + def get_auth_lp(cls) -> bool: + return cls._viva_alp + + @classmethod + def set_auth_lp(cls, value: bool) -> None: + cls._viva_alp = value + @classmethod async def disconnect(cls) -> None: cls._is_connected = False @@ -126,20 +138,37 @@ async def _ear(cls) -> None: NetCode.REQ_LOG_WARNING, NetCode.REQ_LOG_ERROR, ): - await cls._log() + cls._log(code, receive_packet) elif code == NetCode.REQ_AUTH: + cls._server_name = receive_packet.get( + "server_name", "Not_Set_Server_Name" + ) await cls._auth() elif code == NetCode.REQ_FILE_DOWNLOAD: - await cls._download_file() - - elif code == NetCode.END_FILE_DOWNLOAD: - cls._is_downloaded = True + await cls._download_file(receive_packet) else: logger.error("Unknown 'code' type from server") + @classmethod + async def send_packet(self, code: NCAnyType, **kwargs: Any) -> None: + payload = {"code": code, **kwargs} + + await self.send_raw(msgpack.packb(payload)) + + @classmethod + async def send_raw(self, data: bytes) -> None: + if self._writer is None: + raise ValueError("StreamWriter is not set") + + self._writer.write(len(data).to_bytes(4, byteorder="big")) + await self._writer.drain() + + self._writer.write(data) + await self._writer.drain() + @classmethod async def _receive_packet(cls) -> Any: data_size_bytes = await cls._reader.readexactly(4) @@ -148,18 +177,50 @@ async def _receive_packet(cls) -> Any: packed_data = await cls._reader.readexactly(data_size) return msgpack.unpackb(packed_data) - # async def receive_file(self, file_path: Path) -> None: - # try: - # with file_path.open("wb") as file: - # while True: - # data_size_bytes = await self._reader.readexactly(4) - # data_size = int.from_bytes(data_size_bytes, "big") - # - # if data_size == 0: - # break - # - # chunk = await self._reader.readexactly(data_size) - # file.write(chunk) - # - # except Exception as e: - # await self.log_error(f"Error receiving file: {e}") + @classmethod + def _log(cls, code: NCLogType, receive_packet: dict) -> None: + msg = receive_packet.get("message", "Not set") + + if code == NetCode.REQ_LOG_DEBUG: + logger.debug(msg) + + elif code == NetCode.REQ_LOG_INFO: + logger.info(msg) + + elif code == NetCode.REQ_LOG_WARNING: + logger.warning(msg) + + elif code == NetCode.REQ_LOG_ERROR: + logger.error(msg) + + else: + logger.warning(f"Unknown code for log: {receive_packet}") + + @classmethod + async def _auth(cls) -> None: + if cls._viva_alp: + await cls.send_packet( + NetCode.ANSWER_AUTH_ALP, login=cls._login, password=cls._password + ) + + else: + await cls.send_packet( + NetCode.ANSWER_AUTH_REG, login=cls._login, password=cls._password + ) + + @classmethod + async def _download_file(cls, receive_packet: dict) -> None: + try: + file_name = receive_packet.get("file_name", None) + chunk = receive_packet.get("chunk", None) + + if not chunk or file_name: + return + + file_path: Path = cls._content_path / cls._server_name / file_name + + with file_path.open("wb") as file: + file.write(chunk) + + except Exception as e: + await logger.error(f"Error receiving file: {e}") diff --git a/DMBotNetwork/side/side_server.py b/DMBotNetwork/side/side_server.py index 608b58d..683072a 100644 --- a/DMBotNetwork/side/side_server.py +++ b/DMBotNetwork/side/side_server.py @@ -5,8 +5,8 @@ from pathlib import Path from typing import Any, Dict, Optional, Set, get_type_hints -from units import ClientUnit -from utils import NetCode, ServerDB +from ..units import ClientUnit +from ..utils import NetCode, ServerDB logger = logging.getLogger("DMBotNetwork.Server") @@ -19,7 +19,7 @@ class Server: _allow_registration: bool = True _timeout: float = 30.0 - _content_file: Path = None + _server_name: str = None _is_online: bool = False _main_server: Optional[asyncio.AbstractServer] = None @@ -80,7 +80,7 @@ async def start( timeout: float = 30.0, allow_registration: bool = True, base_access_flags: Dict[str, bool] = {}, - content_file: Path | str = None, + server_name: str = "dev_bot", ) -> None: if cls._is_online: logger.warning("Server is already working") @@ -95,13 +95,6 @@ async def start( cls._allow_registration = allow_registration cls._timeout = timeout - content_file = Path(content_file) - if not content_file.is_file(): - logger.warning(f"Server content is invalid: {content_file}") - return - - cls._content_file = content_file - cls._main_server = await asyncio.start_server( cls._main_client_handler, host, main_port ) @@ -216,7 +209,7 @@ async def _main_client_handler( # Auth @classmethod async def _auth(cls, cl_unit: ClientUnit) -> None: - await cl_unit.send_packet(NetCode.REQ_AUTH) + await cl_unit.send_packet(NetCode.REQ_AUTH, server_name=cls._server_name) receive_packet = await asyncio.wait_for(cl_unit._receive_packet(), cls._timeout) if not isinstance(receive_packet, dict): diff --git a/DMBotNetwork/units/client_unit.py b/DMBotNetwork/units/client_unit.py index d980295..4544d9a 100644 --- a/DMBotNetwork/units/client_unit.py +++ b/DMBotNetwork/units/client_unit.py @@ -3,7 +3,8 @@ from typing import Any import msgpack -from utils import NCAnyType, NetCode + +from ..utils import NCAnyType, NetCode class ClientUnit: @@ -45,7 +46,7 @@ def reader(self) -> StreamWriter: # Net async def req_net(self, type: str, **kwargs: Any) -> None: await self.send_packet(NetCode.REQ_NET, type=type, **kwargs) - + # Logs async def log_debug(self, message: str) -> None: await self.send_packet(NetCode.REQ_LOG_DEBUG, message=message) @@ -76,7 +77,9 @@ async def send_raw(self, data: bytes) -> None: await self._writer.drain() # File send - async def send_file(self, file_path: Path, chunk_size: int = 4096) -> None: + async def send_file( + self, file_path: Path, file_name: str, chunk_size: int = 4096 + ) -> None: if self._writer is None: raise ValueError("StreamWriter is not set") @@ -85,10 +88,11 @@ async def send_file(self, file_path: Path, chunk_size: int = 4096) -> None: while True: chunk = file.read(chunk_size) if not chunk: - await self.send_packet(NetCode.END_FILE_DOWNLOAD) break - await self.send_packet(NetCode.REQ_FILE_DOWNLOAD, chunk) + await self.send_packet( + NetCode.REQ_FILE_DOWNLOAD, file_name=file_name, chunk=chunk + ) except Exception as e: await self.log_error(f"Error sending file: {e}") diff --git a/DMBotNetwork/utils/server_db.py b/DMBotNetwork/utils/server_db.py index f5e19d8..f5c29af 100644 --- a/DMBotNetwork/utils/server_db.py +++ b/DMBotNetwork/utils/server_db.py @@ -192,12 +192,14 @@ async def delete_user(cls, username: str) -> None: "DELETE FROM users WHERE username = ?", (username,) ) await cls._connection.commit() - del cls._access_cache[username] + + cls._access_cache.pop(username, None) cls._exist_user.discard(username) except Exception as e: logger.error(f"Error deleting user {username}: {e}") + @classmethod async def change_user_password(cls, username: str, new_password: str) -> None: if username not in cls._exist_user: diff --git a/Tests/ServerDB.py b/Tests/ServerDB.py new file mode 100644 index 0000000..945d9e7 --- /dev/null +++ b/Tests/ServerDB.py @@ -0,0 +1,50 @@ +import unittest +from pathlib import Path +from DMBotNetwork.utils.server_db import ServerDB + +class TestServerDB(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.temp_db_file: Path = Path("temp") + self.temp_db_file.mkdir(exist_ok=True) + ServerDB.set_db_path(self.temp_db_file) + await ServerDB.start() + + async def asyncTearDown(self): + await ServerDB.stop() + db_file = self.temp_db_file / "server.db" + if db_file.exists(): + db_file.unlink() + if self.temp_db_file.exists(): + self.temp_db_file.rmdir() + + async def test_add_user(self): + await ServerDB.add_user("test_user", "password") + self.assertIn("test_user", ServerDB._exist_user) + await ServerDB.delete_user("test_user") + + async def test_delete_user(self): + await ServerDB.add_user("test_user", "password") + await ServerDB.delete_user("test_user") + self.assertNotIn("test_user", ServerDB._exist_user) + + async def test_change_user_password(self): + await ServerDB.add_user("test_user", "password") + await ServerDB.change_user_password("test_user", "new_password") + login = await ServerDB.login_user('test_user', 'new_password') + self.assertEqual(login, 'test_user') + await ServerDB.delete_user("test_user") + + async def test_get_access(self): + await ServerDB.add_user("test_user", "password", {"some_access": True}) + access = await ServerDB.get_access("test_user") + self.assertEqual(access, {"some_access": True}) + await ServerDB.delete_user("test_user") + + async def test_login_user(self): + await ServerDB.add_user("test_user", "password") + result = await ServerDB.login_user("test_user", "password") + self.assertEqual(result, "test_user") + await ServerDB.delete_user("test_user") + +if __name__ == "__main__": + unittest.main() diff --git a/Tests/__init__.py b/Tests/__init__.py index 3724b49..e69de29 100644 --- a/Tests/__init__.py +++ b/Tests/__init__.py @@ -1 +0,0 @@ -# bruh diff --git a/Tests/bruh.py b/Tests/bruh.py deleted file mode 100644 index 3724b49..0000000 --- a/Tests/bruh.py +++ /dev/null @@ -1 +0,0 @@ -# bruh diff --git a/setup.py b/setup.py index fac62a5..270f6bb 100644 --- a/setup.py +++ b/setup.py @@ -4,7 +4,7 @@ name="DMBotNetwork", version="0.1.0", packages=find_packages(), - install_requires=['aiosqlite', 'bcrypt', 'msgpack'], + install_requires=["aiosqlite", "bcrypt", "msgpack"], author="Angels And Demons dev team", author_email="dm.bot.adm@gmail.com", description="Нэткод для проектов DMBot", @@ -16,6 +16,6 @@ "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", "Operating System :: OS Independent", ], - python_requires='>=3.11', + python_requires=">=3.11", license="GPL-3.0", )