diff --git a/include/__init__.py b/include/__init__.py new file mode 100644 index 0000000..00d68ac --- /dev/null +++ b/include/__init__.py @@ -0,0 +1,6 @@ +""" +CFMS WebSocket Server - Core Include Package + +This package contains the core functionality for the CFMS WebSocket server, +including connection handling, database models, request handlers, and utilities. +""" diff --git a/include/classes/__init__.py b/include/classes/__init__.py new file mode 100644 index 0000000..97d32c7 --- /dev/null +++ b/include/classes/__init__.py @@ -0,0 +1,5 @@ +""" +CFMS Classes Module + +Core classes for handling connections, requests, authentication, and access rules. +""" diff --git a/include/classes/auth.py b/include/classes/auth.py index 9501c1f..3264dbd 100644 --- a/include/classes/auth.py +++ b/include/classes/auth.py @@ -1,5 +1,7 @@ +import time from typing import Optional -import jwt, time + +import jwt __all__ = ["Token"] diff --git a/include/classes/connection.py b/include/classes/connection.py index b5918d5..eff8e5e 100644 --- a/include/classes/connection.py +++ b/include/classes/connection.py @@ -1,29 +1,29 @@ +import base64 +import hashlib import json +import mmap +import os import sys import time -import os -import base64 -import hashlib import traceback -import jsonschema from typing import Iterable +from typing import Optional +import jsonschema import websockets -from websockets.sync.server import ServerConnection +from Crypto.Cipher import AES +from Crypto.Random import get_random_bytes from websockets.asyncio.server import broadcast +from websockets.sync.server import ServerConnection from websockets.typing import Data from include.conf_loader import global_config +from include.constants import FILE_TRANSFER_CHUNK_SIZE from include.database.handler import Session from include.database.models.classic import User from include.database.models.file import File, FileTask -from include.util.log import getCustomLogger - from include.shared import connected_listeners - -from Crypto.Cipher import AES -from Crypto.Random import get_random_bytes -import mmap +from include.util.log import getCustomLogger logger = getCustomLogger( "connection", @@ -53,16 +53,20 @@ def __init__(self, websocket: ServerConnection, message: Data) -> None: self.username: str = self.request.get("username", "") self.token: str = self.request.get("token", "") - def conclude_request(self, code: int, data: dict = {}, message: str = "") -> None: + def conclude_request( + self, code: int, data: Optional[dict] = None, message: str = "" + ) -> None: """ Conclude the request by sending a response back to the client. Args: - message: The data/message received from the client. + code: HTTP status code for the response. + data: Data dictionary to include in the response. + message: Message string to include in the response. """ response = { "code": code, - "data": data, + "data": data if data is not None else {}, "message": message, "timestamp": time.time(), } @@ -72,23 +76,6 @@ def conclude_request(self, code: int, data: dict = {}, message: str = "") -> Non self.websocket.send(response_json) - # def authenticate_user(self, user: User|None) -> bool: - # """ - # Authenticates the user by checking the user authentication status. - # Returns: - # bool: True if the user is authenticated, False otherwise. If the user is not authenticated, - # it concludes the request with a 403 status code and an error message indicating - # an invalid user or token. - # """ - - # if not user or not user.is_token_valid(self.token): - # self.conclude_request( - # **{"code": 403, "message": "Invalid user or token", "data": {}} - # ) - # return False - - # return True - def send_file(self, task_id: str) -> None: """ Sends a file associated with the given task ID to the client over a websocket connection using AES encryption. @@ -318,7 +305,7 @@ def receive_file(self, task_id: str) -> None: data = self.websocket.recv() f.write(data) # type: ignore - if not data or len(data) < 8192: + if not data or len(data) < FILE_TRANSFER_CHUNK_SIZE: break except ( websockets.ConnectionClosed, diff --git a/include/conf_loader.py b/include/conf_loader.py index d69c063..ed8ea2a 100644 --- a/include/conf_loader.py +++ b/include/conf_loader.py @@ -1,6 +1,9 @@ -import os, tomllib +import os import secrets -from tomlkit import parse, dumps +import tomllib + +from tomlkit import dumps +from tomlkit import parse __all__ = ["global_config"] diff --git a/include/constants.py b/include/constants.py index 1833187..0493705 100644 --- a/include/constants.py +++ b/include/constants.py @@ -1,9 +1,28 @@ from include.classes.version import Version -# __all__ = ["PROTOCOL_VERSION"] +__all__ = [ + "CORE_VERSION", + "PROTOCOL_VERSION", + "AVAILABLE_ACCESS_TYPES", + "AVAILABLE_BLOCK_TYPES", + "DEFAULT_TOKEN_EXPIRY_SECONDS", + "FAILED_LOGIN_DELAY_SECONDS", + "DEFAULT_SSL_CERT_VALIDITY_DAYS", + "FILE_TRANSFER_CHUNK_SIZE", + "FILE_TASK_DEFAULT_DURATION_SECONDS", +] CORE_VERSION = Version("0.1.0.250919_alpha") PROTOCOL_VERSION = 3 AVAILABLE_ACCESS_TYPES = ["read", "write", "move", "manage"] -AVAILABLE_BLOCK_TYPES: set = {"read", "write", "move"} \ No newline at end of file +AVAILABLE_BLOCK_TYPES: set = {"read", "write", "move"} + +# Authentication and Security Constants +DEFAULT_TOKEN_EXPIRY_SECONDS = 3600 # 1 hour +FAILED_LOGIN_DELAY_SECONDS = 3 # Delay after failed login attempt +DEFAULT_SSL_CERT_VALIDITY_DAYS = 365 # 1 year + +# File Transfer Constants +FILE_TRANSFER_CHUNK_SIZE = 8192 # 8KB - size threshold for determining end of transfer +FILE_TASK_DEFAULT_DURATION_SECONDS = 3600 # 1 hour \ No newline at end of file diff --git a/include/database/__init__.py b/include/database/__init__.py new file mode 100644 index 0000000..6e2585d --- /dev/null +++ b/include/database/__init__.py @@ -0,0 +1,5 @@ +""" +CFMS Database Module + +Database connection handler and ORM models for the CFMS system. +""" diff --git a/include/database/handler.py b/include/database/handler.py index e41ffeb..7e1e7bb 100644 --- a/include/database/handler.py +++ b/include/database/handler.py @@ -1,8 +1,9 @@ -from sqlalchemy import create_engine, URL, MetaData +from sqlalchemy import URL, create_engine +from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import sessionmaker -from include.conf_loader import global_config -from sqlalchemy.orm import DeclarativeBase +from include.conf_loader import global_config +from include.constants import DEFAULT_TOKEN_EXPIRY_SECONDS __all__ = ["engine", "Session", "Base"] @@ -39,13 +40,12 @@ ) engine = create_engine( url, - pool_recycle=3600, + pool_recycle=DEFAULT_TOKEN_EXPIRY_SECONDS, echo=debug_enabled, ) Session = sessionmaker(bind=engine) -# metadata_obj = MetaData() + class Base(DeclarativeBase): - # metadata = metadata_obj pass \ No newline at end of file diff --git a/include/database/models/__init__.py b/include/database/models/__init__.py new file mode 100644 index 0000000..c55fb67 --- /dev/null +++ b/include/database/models/__init__.py @@ -0,0 +1,5 @@ +""" +CFMS Database Models + +ORM models for users, groups, documents, files, and access control. +""" diff --git a/include/database/models/classic.py b/include/database/models/classic.py index 31012cd..668c0ed 100644 --- a/include/database/models/classic.py +++ b/include/database/models/classic.py @@ -1,24 +1,24 @@ +import hashlib +import jwt +import os +import secrets +import time from typing import TYPE_CHECKING from typing import List from typing import Optional from typing import Set -import secrets -from sqlalchemy import VARCHAR, Float, ForeignKey, Integer, Text -from include.database.handler import Base, Session -from include.conf_loader import global_config -from include.classes.auth import Token +from sqlalchemy import VARCHAR, Boolean, Float, ForeignKey, Integer, JSON, Text +from sqlalchemy import event from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column from sqlalchemy.orm import relationship -from sqlalchemy import Boolean -from sqlalchemy import event -import time from sqlalchemy.orm.session import object_session -from sqlalchemy import JSON -import jwt -import hashlib -import os + +from include.classes.auth import Token +from include.conf_loader import global_config +from include.constants import DEFAULT_TOKEN_EXPIRY_SECONDS +from include.database.handler import Base, Session if TYPE_CHECKING: @@ -73,7 +73,7 @@ def authenticate_and_create_token(self, plain_password: str) -> Optional[Token]: else self.secret_key ) token = Token(secret, self.username) - token.new(3600) + token.new(DEFAULT_TOKEN_EXPIRY_SECONDS) session = object_session(self) if session is not None: @@ -113,7 +113,7 @@ def renew_token(self) -> Token: else self.secret_key ) new_token = Token(secret, self.username) - new_token.new(3600) + new_token.new(DEFAULT_TOKEN_EXPIRY_SECONDS) return new_token diff --git a/include/database/models/file.py b/include/database/models/file.py index 910bbec..449924e 100644 --- a/include/database/models/file.py +++ b/include/database/models/file.py @@ -1,14 +1,17 @@ +import os import secrets -from sqlalchemy import VARCHAR, Float, ForeignKey, Integer, Text, Boolean -from include.database.handler import Base +import sys +import time from typing import List from typing import Optional + +from sqlalchemy import VARCHAR, Float, ForeignKey, Integer, Text, Boolean from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column from sqlalchemy.orm import relationship -import time from sqlalchemy.orm.session import object_session -import os, sys + +from include.database.handler import Base class File(Base): diff --git a/include/handlers/__init__.py b/include/handlers/__init__.py new file mode 100644 index 0000000..00bcf79 --- /dev/null +++ b/include/handlers/__init__.py @@ -0,0 +1,6 @@ +""" +CFMS Request Handlers + +Request handlers for authentication, document management, directory operations, +and system management. +""" diff --git a/include/handlers/auth.py b/include/handlers/auth.py index d4e895b..b679b30 100644 --- a/include/handlers/auth.py +++ b/include/handlers/auth.py @@ -1,11 +1,12 @@ +import time + from include.classes.connection import ConnectionHandler from include.classes.request import RequestHandler from include.conf_loader import global_config +from include.constants import FAILED_LOGIN_DELAY_SECONDS from include.database.handler import Session from include.database.models.classic import User from include.util.audit import log_audit -import time - from include.util.pwd import check_passwd_requirements @@ -99,7 +100,7 @@ def handle(self, handler: ConnectionHandler): response = response_invalid if response == response_invalid: - time.sleep(3) + time.sleep(FAILED_LOGIN_DELAY_SECONDS) # Send the response back to the client handler.conclude_request(**response) diff --git a/include/handlers/document.py b/include/handlers/document.py index ea43bef..6367364 100644 --- a/include/handlers/document.py +++ b/include/handlers/document.py @@ -1,26 +1,26 @@ import datetime import secrets +import time import jsonschema + from include.classes.connection import ConnectionHandler from include.classes.request import RequestHandler +from include.conf_loader import global_config from include.constants import AVAILABLE_ACCESS_TYPES +from include.constants import FILE_TASK_DEFAULT_DURATION_SECONDS from include.database.handler import Session -from include.database.models.classic import ( - User, -) +from include.database.models.classic import User from include.database.models.entity import ( Document, - Folder, - NoActiveRevisionsError, DocumentAccessRule, DocumentRevision, + Folder, + NoActiveRevisionsError, ) from include.database.models.file import File, FileTask -from include.conf_loader import global_config from include.util.rule.applying import apply_access_rules import include.system.messages as smsg -import time __all__ = [ "RequestGetDocumentInfoHandler", @@ -36,8 +36,7 @@ ] -# def create_file_task(file_id: str, transfer_mode=0): -def create_file_task(file: File, transfer_mode=0): +def create_file_task(file: File, transfer_mode: int = 0): """ Creates a new file processing task for the specified file. Args: @@ -52,17 +51,17 @@ def create_file_task(file: File, transfer_mode=0): """ with Session() as session: - # file = session.get(File, file_id) if not file: return None + now = time.time() task = FileTask( file_id=file.id, status=0, mode=transfer_mode, start_time=now, - end_time=now + 3600, + end_time=now + FILE_TASK_DEFAULT_DURATION_SECONDS, ) session.add(task) session.commit() diff --git a/include/handlers/management/__init__.py b/include/handlers/management/__init__.py new file mode 100644 index 0000000..b1dc4f6 --- /dev/null +++ b/include/handlers/management/__init__.py @@ -0,0 +1,6 @@ +""" +CFMS Management Handlers + +Request handlers for user management, group management, access control, +and system administration. +""" diff --git a/include/system/__init__.py b/include/system/__init__.py new file mode 100644 index 0000000..0da3283 --- /dev/null +++ b/include/system/__init__.py @@ -0,0 +1,5 @@ +""" +CFMS System Module + +System-level messages and utilities for server operations. +""" diff --git a/include/util/__init__.py b/include/util/__init__.py new file mode 100644 index 0000000..ffefd75 --- /dev/null +++ b/include/util/__init__.py @@ -0,0 +1,6 @@ +""" +CFMS Utilities + +Utility functions for logging, user management, group management, +password validation, and audit logging. +""" diff --git a/include/util/audit.py b/include/util/audit.py index 4bb3ce1..9911bda 100644 --- a/include/util/audit.py +++ b/include/util/audit.py @@ -1,4 +1,5 @@ from typing import Optional + from include.database.handler import Session from include.database.models.classic import User, AuditEntry @@ -10,7 +11,7 @@ def log_audit( target: Optional[str] = None, data: Optional[dict] = None, remote_address: Optional[str] = None, -): +) -> None: """创建审计日志。""" if result == 400: return diff --git a/include/util/log.py b/include/util/log.py index fdf15de..d7120f1 100644 --- a/include/util/log.py +++ b/include/util/log.py @@ -1,5 +1,6 @@ import logging from logging.handlers import RotatingFileHandler +from typing import Tuple """ Provides a utility util to create and configure a custom logger with both file and console handlers. @@ -18,8 +19,10 @@ def getCustomLogger( - logname: str, level: tuple=(logging.DEBUG, logging.INFO), filepath="default.log" -): + logname: str, + level: Tuple[int, int] = (logging.DEBUG, logging.INFO), + filepath: str = "default.log" +) -> logging.Logger: logger = logging.getLogger(logname) logger.setLevel(level=logging.DEBUG) # This level must be 'logging.DEBUG'. logger.propagate = False diff --git a/include/util/pwd.py b/include/util/pwd.py index 00877dd..b360002 100644 --- a/include/util/pwd.py +++ b/include/util/pwd.py @@ -1,3 +1,4 @@ +from typing import List from typing import Optional __all__ = [ @@ -8,10 +9,10 @@ class MissingComponentsError(ValueError): - def __init__(self, missing: set[str]): + def __init__(self, missing: set[str]) -> None: self.missing = missing - def __str__(self): + def __str__(self) -> str: return ( f"Password is missing the necessary characters: {", ".join(self.missing)}" ) @@ -23,7 +24,7 @@ def __init__( length: int, min_length: Optional[int] = None, max_length: Optional[int] = None, - ): + ) -> None: self.length = length self.min_length = min_length self.max_length = max_length @@ -31,7 +32,7 @@ def __init__( not self.min_length and not self.max_length ) - def __str__(self): + def __str__(self) -> str: if self.min_length and self.max_length: return f"Password does not meet the length requirement ({self.min_length} ~ {self.max_length})" @@ -43,13 +44,16 @@ def check_passwd_requirements( passwd: str, min_length: int, max_length: int, - must_contain: list[list[str]] = [], -): + must_contain: Optional[List[List[str]]] = None, +) -> None: length = len(passwd) if not (min_length <= length <= max_length): raise InvaildPasswordLengthError(length, min_length, max_length) pwd_set = set(passwd) + + if must_contain is None: + must_contain = [] for group in must_contain: each_set = set(group) diff --git a/include/util/rule/__init__.py b/include/util/rule/__init__.py new file mode 100644 index 0000000..0773244 --- /dev/null +++ b/include/util/rule/__init__.py @@ -0,0 +1,5 @@ +""" +CFMS Access Rule Utilities + +Utilities for validating and applying access control rules. +""" diff --git a/include/util/user.py b/include/util/user.py index fc37684..4d4afe3 100644 --- a/include/util/user.py +++ b/include/util/user.py @@ -1,5 +1,8 @@ -import string, secrets, hashlib +import hashlib +import secrets +import string import time + from include.database.models.classic import User, UserMembership, UserPermission from include.database.handler import Session diff --git a/main.py b/main.py index bf86efd..0c7c439 100644 --- a/main.py +++ b/main.py @@ -26,18 +26,22 @@ os.makedirs("./content/logs/", exist_ok=True) os.makedirs("./content/ssl/", exist_ok=True) +import socket import ssl + +from websockets.sync.server import serve + from include.conf_loader import global_config +from include.connection_handler import handle_connection from include.constants import CORE_VERSION -from include.database.handler import engine, Base +from include.constants import DEFAULT_SSL_CERT_VALIDITY_DAYS +from include.database.handler import Base from include.database.handler import Session +from include.database.handler import engine from include.database.models.classic import User, UserGroup from include.database.models.entity import Document, DocumentRevision from include.database.models.file import File -from websockets.sync.server import serve -from include.connection_handler import handle_connection from include.util.log import getCustomLogger -import socket def server_init(): @@ -175,7 +179,7 @@ def server_init(): .not_valid_before(datetime.datetime.now(datetime.timezone.utc)) .not_valid_after( datetime.datetime.now(datetime.timezone.utc) - + datetime.timedelta(days=365) + + datetime.timedelta(days=DEFAULT_SSL_CERT_VALIDITY_DAYS) ) .add_extension( x509.SubjectAlternativeName( diff --git a/test.py b/test.py index 8e997a6..7429067 100644 --- a/test.py +++ b/test.py @@ -2,8 +2,10 @@ """Client using the threading API.""" +import json +import ssl + from websockets.sync.client import connect -import ssl, json print("Hello world! Client")