diff --git a/burrito/__init__.py b/burrito/__init__.py index 854207e1..247ae08c 100644 --- a/burrito/__init__.py +++ b/burrito/__init__.py @@ -1,6 +1,6 @@ from pytz import timezone -__version__ = "0.8.0 indev" +__version__ = "1.0.0 alpha" CURRENT_TIME_ZONE = timezone("Europe/Kyiv") diff --git a/burrito/apps/comments/views.py b/burrito/apps/comments/views.py index 81a745d8..fc8676f6 100644 --- a/burrito/apps/comments/views.py +++ b/burrito/apps/comments/views.py @@ -14,7 +14,13 @@ from burrito.models.tickets_model import Tickets from burrito.utils.mongo_util import mongo_insert, mongo_update, mongo_delete -from burrito.utils.tickets_util import is_ticket_exist, am_i_own_this_ticket, send_notification, make_short_user_data +from burrito.utils.tickets_util import ( + is_ticket_exist, + am_i_own_this_ticket, + send_notification, + make_short_user_data, + send_comment_update +) from burrito.utils.permissions_checker import check_permission from burrito.utils.auth import get_auth_core, AuthTokenPayload, BurritoJWT from burrito.utils.query_util import STATUS_OPEN @@ -61,6 +67,7 @@ async def comments__create( body=f"Someone has created a new comment in ticket {ticket.ticket_id}" ) ) + send_comment_update(ticket.ticket_id, comment_id) if ticket.status.status_id in (4, 6) and am_i_own_this_ticket(ticket.creator.user_id, token_payload.user_id): create_ticket_action( diff --git a/burrito/apps/ws/utils.py b/burrito/apps/ws/utils.py index 917fdc0c..a6969d67 100644 --- a/burrito/apps/ws/utils.py +++ b/burrito/apps/ws/utils.py @@ -1,15 +1,19 @@ import time import threading +from fastapi import HTTPException from redis.client import PubSub from websockets.sync.server import serve from websockets.legacy.server import WebSocketServerProtocol +from burrito.models.tickets_model import Tickets + from burrito.utils.config_reader import get_config from burrito.utils.logger import get_logger from burrito.utils.auth import check_jwt_token from burrito.utils.redis_utils import get_redis_connector from burrito.utils.auth import AuthTokenPayload +from burrito.utils.tickets_util import is_ticket_exist def recv_data(websocket: WebSocketServerProtocol) -> bytes: @@ -36,45 +40,77 @@ def recv_ping(websocket: WebSocketServerProtocol) -> None: if raw_data == "PING": send_data(websocket, b"PONG") - get_logger().info("PING/PONG") except Exception as exc: - get_logger().warning(f"PING/PONG {exc}") + get_logger().warning(f"PING/PONG {exc}", exc_info=True) break time.sleep(5) get_logger().info(f"Pinging thread is finished {thread_id}") -def main_handler(websocket: WebSocketServerProtocol): - thread_id = threading.get_native_id() - get_logger().info(f"New thread started: {thread_id}") +def notifications_cycle(websocket: WebSocketServerProtocol, token_payload: AuthTokenPayload): + pubsub: PubSub = get_redis_connector().pubsub() + pubsub.subscribe(f"user_{token_payload.user_id}") - raw_data = recv_data(websocket) + pinging_thread = threading.Thread(target=recv_ping, args=(websocket,), daemon=True) + pinging_thread.start() - if isinstance(raw_data, bytes): - raw_data = raw_data.decode("utf-8") + try: + while True: + if not pinging_thread.is_alive(): + break - if not raw_data: - send_data(websocket, b"Auth fail") + message = pubsub.get_message() - token_payload: AuthTokenPayload = None - try: - token_payload = check_jwt_token(raw_data) + if message: + data = message.get("data") - except: - send_data(websocket, b"Auth fail") + if isinstance(data, bytes): + send_data(websocket, data) + + time.sleep(0.5) + + except Exception as exc: + get_logger().warning(f"{exc}", exc_info=True) + + +def chat_cycle(websocket: WebSocketServerProtocol, token_payload: AuthTokenPayload): + chat_number = recv_data(websocket) + if chat_number.isdigit(): + chat_number = int(chat_number) + + try: + ticket: Tickets = is_ticket_exist(chat_number) + + if ticket.hidden and token_payload.user_id not in ( + ticket.creator.user_id, + ticket.assignee.user_id if ticket.assignee else -1 + ): + send_data(websocket, b"Is not allowed to interact with this ticket") + close_conn(websocket) + return + + except HTTPException: + get_logger().warning(f"Ticket with ID {chat_number} is not exist") + send_data(websocket, f"Ticket with ID {chat_number} is not exist") close_conn(websocket) return - send_data(websocket, b"Auth OK") + except Exception: + get_logger().warning(f"User {token_payload.user_id} tried to join chat {ticket.ticket_id}", exc_info=True) + close_conn(websocket) + + get_logger().info(f"User {token_payload.user_id} joined the chat ({chat_number})") pubsub: PubSub = get_redis_connector().pubsub() - pubsub.subscribe(f"user_{token_payload.user_id}") + pubsub.subscribe(f"chat_{chat_number}") pinging_thread = threading.Thread(target=recv_ping, args=(websocket,), daemon=True) pinging_thread.start() + send_data(websocket, f"Successfully joined to chat {chat_number}".encode("utf-8")) + try: while True: if not pinging_thread.is_alive(): @@ -91,7 +127,53 @@ def main_handler(websocket: WebSocketServerProtocol): time.sleep(0.5) except Exception as exc: - get_logger().warning(f"{exc}") + get_logger().warning(f"{exc}", exc_info=True) + + get_logger().info(f"User {token_payload.user_id} left the chat ({chat_number})") + + +__CONNECTION_MODES = { + "NOTIFICATIONS": notifications_cycle, + "CHAT": chat_cycle +} + + +def main_handler(websocket: WebSocketServerProtocol): + thread_id = threading.get_native_id() + get_logger().info(f"New thread started: {thread_id}") + + raw_data = recv_data(websocket) + + if isinstance(raw_data, bytes): + raw_data = raw_data.decode("utf-8") + + if not raw_data: + send_data(websocket, b"Auth fail") + + token_payload: AuthTokenPayload = None + try: + token_payload = check_jwt_token(raw_data) + + except Exception: + send_data(websocket, b"Auth fail") + close_conn(websocket) + return + + send_data(websocket, b"Auth OK") + + connection_mode = recv_data(websocket) + if isinstance(connection_mode, bytes): + connection_mode = connection_mode.decode("utf-8") + if not connection_mode: + send_data(websocket, b"Connection mode is invalid") + + connection_handler = __CONNECTION_MODES.get(connection_mode) + if connection_handler: + send_data(websocket, b"Connection mode is OK") + connection_handler(websocket, token_payload) + else: + send_data(websocket, b"Connection mode is invalid") + get_logger().warning(f"User {token_payload.user_id} has selected wrong connection mode '{connection_mode}'") get_logger().info(f"Thread {thread_id} is finished") diff --git a/burrito/utils/hash_util.py b/burrito/utils/hash_util.py index c1bbffe5..742b3a63 100644 --- a/burrito/utils/hash_util.py +++ b/burrito/utils/hash_util.py @@ -36,17 +36,5 @@ def compare_password(password: str, hashed_password: str) -> bool: try: return _hasher.verify(hashed_password, password) - except: + except Exception: return False - - -def get_verification_code() -> str: - """_summary_ - - Generate and return email verification code - - Returns: - str: verification code - """ - - return "".join((str(SystemRandom().randint(0, 9)) for i in range(6))) diff --git a/burrito/utils/tickets_util.py b/burrito/utils/tickets_util.py index b15fab9c..db614dc7 100644 --- a/burrito/utils/tickets_util.py +++ b/burrito/utils/tickets_util.py @@ -431,3 +431,7 @@ def send_notification(ticket: Tickets | int, notification: Notifications): if mongo_items_count(NotificationMetaData, notification_id=notification_id) == 0: mongo_delete(Notifications, _id=notification_id) + + +def send_comment_update(ticket_id: int, comment_id: str) -> None: + get_redis_connector().publish(f"chat_{ticket_id}", comment_id.encode("utf-8")) diff --git a/burrito/utils/websockets.py b/burrito/utils/websockets.py index 2fb03fc0..1bfb708e 100644 --- a/burrito/utils/websockets.py +++ b/burrito/utils/websockets.py @@ -1,49 +1,9 @@ import orjson as json -from starlette.websockets import WebSocket - -from burrito.utils.singleton_pattern import singleton -from burrito.utils.auth import check_jwt_token from burrito.models.ws_message import WebSocketMessage from burrito.models.m_notifications_model import Notifications -@singleton -class WebsocketManager: - def __init__(self) -> None: - self._websocket_list: list[WebSocket] = [] - - async def accept(self, websocket: WebSocket): - await websocket.accept() - self._websocket_list.append(websocket) - - async def check_token(self, websocket: WebSocket) -> bool: - token = (await self.recv(websocket)).decode("utf-8") - - try: - return await check_jwt_token(token) - except: - await self.send(websocket, b"Auth fail") - await self.close(websocket) - - async def recv(self, websocket: WebSocket) -> bytes: - return await websocket.receive_bytes() - - async def send(self, websocket: WebSocket, data: bytes): - await websocket.send_bytes(data) - - async def close(self, websocket: WebSocket): - self._websocket_list.remove(websocket) - await websocket.close() - - -def get_websocket_manager(): - return WebsocketManager() - - -_WEBSOCKET_MANAGER: WebsocketManager = get_websocket_manager() - - def make_websocket_message(type_: str, obj: Notifications) -> bytes: return json.dumps( WebSocketMessage( diff --git a/pyproject.toml b/pyproject.toml index 8193482f..2832d0e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "Burrito" -version = "0.8.0.dev2" +version = "1.0.0.alpha" description = "API for the issue tracker" authors = ["DimonBor", "m-o-d-e-r"] readme = "README.md" diff --git a/scripts/burrito_cluster_ping.py b/scripts/burrito_cluster_ping.py index 7f46bb46..a56c032a 100755 --- a/scripts/burrito_cluster_ping.py +++ b/scripts/burrito_cluster_ping.py @@ -1,6 +1,5 @@ import socket -import requests from dotenv import dotenv_values, find_dotenv import rich @@ -26,7 +25,7 @@ def _ping(data): try: sock.connect(("localhost", int(data[1]))) rich.print(f"\t[green][+] Host is OK: {data}") - except: + except Exception: rich.print(f"\t[red][-] No connection to host: {data}") finally: sock.close()