refactor: ruff autofixes

This commit is contained in:
Mark Paronyan committed 2026-09-27 15:55:10 +03:00
1 parent 2ac637c662
commit a0450b6bce
455 files changed
+12322 -14762

No files matched your search

+1
View File
@@ -111,6 +111,7 @@ def is_ipv6_enabled() -> bool:
return True
raise
def md5sum(filename):
"""
Calculate the MD5 checksum of a file.
+3 -1
View File
@@ -73,7 +73,9 @@ async def subprocess_check_output(*args, cwd=None, env=None, stderr=False):
proc = await asyncio.create_subprocess_exec(*args, stderr=asyncio.subprocess.PIPE, cwd=cwd, env=env)
_, output = await proc.communicate()
else:
proc = await asyncio.create_subprocess_exec(*args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL, cwd=cwd, env=env)
proc = await asyncio.create_subprocess_exec(
*args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL, cwd=cwd, env=env
)
output, _ = await proc.communicate()
if output is None:
return ""
+14 -14
View File
@@ -56,20 +56,24 @@ def _get_compressor(compress_type, compresslevel=None):
if compress_type == zipfile.ZIP_DEFLATED:
from zipfile import zlib
if compresslevel is not None:
return zlib.compressobj(compresslevel, zlib.DEFLATED, -15)
return zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -15)
elif compress_type == zipfile.ZIP_BZIP2:
from zipfile import bz2
if compresslevel is not None:
return bz2.BZ2Compressor(compresslevel)
return bz2.BZ2Compressor()
# compresslevel is ignored for ZIP_LZMA
elif compress_type == zipfile.ZIP_LZMA:
from zipfile import LZMACompressor
return LZMACompressor()
elif compress_type == ZIP_ZSTANDARD:
import zstandard as zstd
if compresslevel is not None:
# for threads param a value of -1 means to set the number of threads to the number of detected logical CPUs
return zstd.ZstdCompressor(level=compresslevel, threads=-1).compressobj()
@@ -141,13 +145,13 @@ class ZipInfo(zipfile.ZipInfo):
class ZipFile(zipfile.ZipFile):
def __init__(
self,
fileobj=None,
mode="w",
compression=zipfile.ZIP_STORED,
allowZip64=True,
compresslevel=None,
chunksize=32768
self,
fileobj=None,
mode="w",
compression=zipfile.ZIP_STORED,
allowZip64=True,
compresslevel=None,
chunksize=32768,
):
"""Open the ZIP file with mode write "w"."""
@@ -158,11 +162,7 @@ class ZipFile(zipfile.ZipFile):
self._comment = b""
zipfile.ZipFile.__init__(
self, fileobj,
mode=mode,
compression=compression,
compresslevel=compresslevel,
allowZip64=allowZip64
self, fileobj, mode=mode, compression=compression, compresslevel=compresslevel, allowZip64=allowZip64
)
self._chunksize = chunksize
self.paths_to_write = []
@@ -232,7 +232,7 @@ class ZipFile(zipfile.ZipFile):
"filename": filename,
"arcname": arcname,
"compress_type": compress_type,
"compresslevel": compresslevel
"compresslevel": compresslevel,
}
self.paths_to_write.append(kwargs)
@@ -245,7 +245,7 @@ class ZipFile(zipfile.ZipFile):
"arcname": arcname,
"iterable": iterable,
"compress_type": compress_type,
"compresslevel": compresslevel
"compresslevel": compresslevel,
}
self.paths_to_write.append(kwargs)
+6 -6
View File
@@ -114,7 +114,7 @@ class EmbedShell:
if not cmd[0].strip():
return ""
for (name, meth) in inspect.getmembers(self):
for name, meth in inspect.getmembers(self):
if name == cmd[0]:
cmd.pop(0)
res = await meth(*cmd)
@@ -177,10 +177,10 @@ class UnstoppableEventLoop(EventLoop):
self._loop = loop
def close(self):
" Ignore. "
"Ignore."
def stop(self):
" Ignore. "
"Ignore."
def run_in_executor(self, *args, **kwargs):
return self._loop.run_in_executor(*args, **kwargs)
@@ -264,18 +264,18 @@ class ShellConnection(TelnetConnection):
self.reset()
def reset(self):
""" Resets terminal screen"""
"""Resets terminal screen"""
self._cli.reset()
self._cli.buffers[DEFAULT_BUFFER].reset()
self._cli.renderer.request_absolute_cursor_position()
self._cli._redraw()
def write(self, data):
""" Compat with CLI"""
"""Compat with CLI"""
self.send(data)
def flush(self):
""" Compat with CLI"""
"""Compat with CLI"""
pass
+1 -2
View File
@@ -4,7 +4,6 @@ Parser for VT100 input stream.
# Copied from prompt_toolkit/terminal/vt100_input.py due to dependency on termios (which is not available on Windows)
import re
import six
@@ -40,7 +39,7 @@ _mouse_event_prefix_re = re.compile("^" + re.escape("\x1b[") + r"(<?[\d;]*|M.{0,
class _Flush:
""" Helper object to indicate flush operation to the parser. """
"""Helper object to indicate flush operation to the parser."""
pass
@@ -48,7 +48,7 @@ class AsyncioRawCommandServer:
*self._command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.STDOUT,
stdin=asyncio.subprocess.PIPE
stdin=asyncio.subprocess.PIPE,
)
try:
await self._process(network_reader, network_writer, process.stdout, process.stdin)
+3 -3
View File
@@ -119,9 +119,9 @@ class AsyncioTelnetServer:
:param keepalive_interval: interval in seconds for sending NOP keep-alive (0 to disable).
"""
assert connection_factory is None or (
connection_factory is not None and reader is None and writer is None
), "Please use either reader and writer either connection_factory, otherwise duplicate data may be produced."
assert connection_factory is None or (connection_factory is not None and reader is None and writer is None), (
"Please use either reader and writer either connection_factory, otherwise duplicate data may be produced."
)
self._reader = reader
self._writer = writer
+8 -23
View File
@@ -42,10 +42,7 @@ def init_encryption(secrets_dir: str) -> None:
encryption_key_path = os.path.join(secrets_dir, "gns3_encryption_key")
if not os.path.exists(encryption_key_path):
log.info(
f"No encryption key found, generating one in "
f"'{encryption_key_path}'..."
)
log.info(f"No encryption key found, generating one in '{encryption_key_path}'...")
try:
key = Fernet.generate_key()
os.makedirs(secrets_dir, exist_ok=True)
@@ -55,26 +52,17 @@ def init_encryption(secrets_dir: str) -> None:
# Set restrictive permissions (owner read/write only)
os.chmod(encryption_key_path, 0o600)
except OSError as e:
log.error(
"Could not create encryption key file "
f"'{encryption_key_path}': {e}"
)
log.error(f"Could not create encryption key file '{encryption_key_path}': {e}")
raise
try:
with open(encryption_key_path, encoding="utf-8") as f:
key_content = f.read().strip()
key_bytes = (
key_content.encode()
if isinstance(key_content, str)
else key_content
)
key_bytes = key_content.encode() if isinstance(key_content, str) else key_content
_fernet = Fernet(key_bytes)
log.debug("Encryption initialized successfully")
except OSError as e:
log.error(
f"Could not read encryption key file '{encryption_key_path}': {e}"
)
log.error(f"Could not read encryption key file '{encryption_key_path}': {e}")
raise
@@ -88,9 +76,7 @@ def encrypt(plaintext: str) -> str:
"""
if _fernet is None:
raise RuntimeError(
"Encryption not initialized. Call init_encryption() first."
)
raise RuntimeError("Encryption not initialized. Call init_encryption() first.")
if not plaintext:
return ""
@@ -110,9 +96,7 @@ def decrypt(ciphertext: str) -> str:
"""
if _fernet is None:
raise RuntimeError(
"Encryption not initialized. Call init_encryption() first."
)
raise RuntimeError("Encryption not initialized. Call init_encryption() first.")
if not ciphertext:
return ""
@@ -139,10 +123,11 @@ def is_encrypted(value: str) -> bool:
# Check for Fernet prefix and valid format
try:
# Fernet tokens always start with 'gAAAAA' in base64
if not value.startswith('gAAAAA'):
if not value.startswith("gAAAAA"):
return False
# Attempt to decode as URL-safe base64
import base64
decoded = base64.urlsafe_b64decode(value)
# Fernet tokens have a specific format (minimum 32 bytes)
return len(decoded) >= 32
+1
View File
@@ -26,6 +26,7 @@ except ImportError:
from contextlib import ExitStack
resource_manager = ExitStack()
atexit.register(resource_manager.close)
+11 -11
View File
@@ -38,15 +38,15 @@ def to_ios_hostname(name):
"""
# Replace invalid characters with hyphens
name = re.sub(r'[^a-zA-Z0-9-]', '-', name)
name = re.sub(r"[^a-zA-Z0-9-]", "-", name)
# Ensure the hostname starts with a letter
if not re.search(r'^[a-zA-Z]', name):
name = 'a' + name
if not re.search(r"^[a-zA-Z]", name):
name = "a" + name
# Ensure the hostname ends with a letter or digit
if not re.search(r'[a-zA-Z0-9]$', name):
name = name.rstrip('-') + '0'
if not re.search(r"[a-zA-Z0-9]$", name):
name = name.rstrip("-") + "0"
# Truncate the hostname to 63 characters
name = name[:63]
@@ -87,24 +87,24 @@ def to_rfc1123_hostname(name: str) -> str:
"""
# Replace invalid characters with hyphens
name = re.sub(r'[^a-zA-Z0-9-.]', '-', name)
name = re.sub(r"[^a-zA-Z0-9-.]", "-", name)
# Remove trailing dot if it exists
name = name.rstrip('.')
name = name.rstrip(".")
# Ensure each label is not longer than 63 characters
labels = name.split('.')
labels = name.split(".")
labels = [label[:63] for label in labels]
# Remove leading and trailing hyphens from each label if they exist
labels = [label.strip('-') for label in labels]
labels = [label.strip("-") for label in labels]
# Check if the TLD is all-numeric and if so, replace it with "invalid"
if re.match(r"[0-9]+$", labels[-1]):
labels[-1] = 'invalid'
labels[-1] = "invalid"
# Join the labels back together
name = '.'.join(labels)
name = ".".join(labels)
# Ensure the total length is not longer than 253 characters
name = name[:253]
+28 -27
View File
@@ -55,7 +55,6 @@ async def list_images(image_type):
default_directory = default_images_directory(image_type)
for directory in images_directories(image_type):
# We limit recursion to path outside the default images directory
# the reason is in the default directory manage file organization and
# it should be flatten to keep things simple
@@ -89,15 +88,19 @@ async def list_images(image_type):
with open(os.path.join(root, filename), "rb") as f:
# read the first 7 bytes of the file.
elf_header_start = f.read(7)
if image_type == "dynamips" and elf_header_start != b'\x7fELF\x01\x02\x01':
if image_type == "dynamips" and elf_header_start != b"\x7fELF\x01\x02\x01":
# IOS images must start with the ELF magic number, be 32-bit, big endian and have an ELF version of 1
log.warning(f"IOS image {filename} does not start with a valid ELF magic number, skipping...")
continue
elif image_type == "iou" and elf_header_start != b'\x7fELF\x02\x01\x01' and elf_header_start != b'\x7fELF\x01\x01\x01':
elif (
image_type == "iou"
and elf_header_start != b"\x7fELF\x02\x01\x01"
and elf_header_start != b"\x7fELF\x01\x01\x01"
):
# IOU images must start with the ELF magic number, be 32-bit or 64-bit, little endian and have an ELF version of 1
log.warning(f"IOU image {filename} does not start with a valid ELF magic number, skipping...")
continue
elif image_type == "qemu" and elf_header_start[:4] == b'\x7fELF':
elif image_type == "qemu" and elf_header_start[:4] == b"\x7fELF":
# QEMU images should not start with an ELF magic number
log.warning(f"QEMU image {filename} starts with an ELF magic number, skipping...")
continue
@@ -117,7 +120,7 @@ async def list_images(image_type):
def get_builtin_disks() -> List[str]:
builtin_disks = []
for entry in importlib_resources.files('gns3server').joinpath("disks").iterdir():
for entry in importlib_resources.files("gns3server").joinpath("disks").iterdir():
if entry.is_file():
builtin_disks.append(entry.name)
return builtin_disks
@@ -132,8 +135,10 @@ async def read_image_info(path: str, expected_image_type: str = None) -> dict:
if len(image_header) >= header_magic_len:
detected_image_type = check_valid_image_header(path, image_header)
if expected_image_type and detected_image_type != expected_image_type:
raise InvalidImageError(f"Detected image type for '{path}' is {detected_image_type}, "
f"expected type is {expected_image_type}")
raise InvalidImageError(
f"Detected image type for '{path}' is {detected_image_type}, "
f"expected type is {expected_image_type}"
)
else:
raise InvalidImageError(f"Image '{path}' is too small to be valid")
except OSError as e:
@@ -306,7 +311,6 @@ def remove_checksum(path):
class InvalidImageError(Exception):
def __init__(self, message: str):
super().__init__()
self._message = message
@@ -317,14 +321,14 @@ class InvalidImageError(Exception):
def check_valid_image_header(path: str, data: bytes, allow_raw_image: bool = False) -> str:
if data[:7] == b'\x7fELF\x01\x02\x01':
if data[:7] == b"\x7fELF\x01\x02\x01":
# for IOS images: file must start with the ELF magic number, be 32-bit, big endian and have an ELF version of 1
return "ios"
elif data[:7] == b'\x7fELF\x01\x01\x01' or data[:7] == b'\x7fELF\x02\x01\x01':
elif data[:7] == b"\x7fELF\x01\x01\x01" or data[:7] == b"\x7fELF\x02\x01\x01":
# for IOU images: file must start with the ELF magic number, be 32-bit or 64-bit, little endian and
# have an ELF version of 1 (normal IOS images are big endian!)
return "iou"
elif data[:4] == b'QFI\xfb' or data[:4] == b'KDMV':
elif data[:4] == b"QFI\xfb" or data[:4] == b"KDMV":
# for Qemy images: file must be QCOW2 or VMDK
return "qemu"
else:
@@ -334,12 +338,12 @@ def check_valid_image_header(path: str, data: bytes, allow_raw_image: bool = Fal
async def write_image(
image_filename: str,
image_path: str,
stream: AsyncGenerator[bytes, None],
images_repo: ImagesRepository,
check_image_header=True,
allow_raw_image=False
image_filename: str,
image_path: str,
stream: AsyncGenerator[bytes, None],
images_repo: ImagesRepository,
check_image_header=True,
allow_raw_image=False,
) -> models.Image:
image_dir, image_name = os.path.split(image_filename)
@@ -370,15 +374,17 @@ async def write_image(
image_path = os.path.abspath(os.path.join(directory, image_filename))
if os.path.exists(image_path):
raise InvalidImageError(f"File '{image_path}' already exists, "
f"please choose a different name or remove the existing image")
raise InvalidImageError(
f"File '{image_path}' already exists, please choose a different name or remove the existing image"
)
checksum = checksum.hexdigest()
image_dir = os.path.dirname(image_path)
duplicate_image = await images_repo.get_image_by_checksum(checksum, image_dir)
if duplicate_image:
raise InvalidImageError(f"Image '{duplicate_image.filename}' with the "
f"same checksum already exists in '{image_dir}'")
raise InvalidImageError(
f"Image '{duplicate_image.filename}' with the same checksum already exists in '{image_dir}'"
)
shutil.move(tmp_path, image_path)
os.chmod(image_path, stat.S_IWRITE | stat.S_IREAD | stat.S_IEXEC)
@@ -390,10 +396,5 @@ async def write_image(
log.warning(f"Could not remove '{tmp_path}'")
return await images_repo.add_image(
image_name,
image_type,
image_size,
image_path,
checksum,
checksum_algorithm="md5"
image_name, image_type, image_size, image_path, checksum, checksum_algorithm="md5"
)
+5 -10
View File
@@ -49,11 +49,11 @@ def _get_windows_interfaces_from_registry():
netcard, _ = winreg.QueryValueEx(hkeycard, "Description")
connection = (
r"SYSTEM\CurrentControlSet\Control\Network\{4D36E972-E325-11CE-BFC1-08002BE10318}"
+ fr"\{guid}\Connection"
+ rf"\{guid}\Connection"
)
hkeycon = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, connection)
name, _ = winreg.QueryValueEx(hkeycon, "Name")
interface = fr"SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces\{guid}"
interface = rf"SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\Interfaces\{guid}"
hkeyinterface = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, interface)
is_dhcp_enabled, _ = winreg.QueryValueEx(hkeyinterface, "EnableDHCP")
if is_dhcp_enabled:
@@ -160,7 +160,6 @@ def is_interface_up(interface):
"""
if sys.platform.startswith("linux") or sys.platform.startswith("openbsd"):
if interface not in psutil.net_if_addrs():
return False
@@ -218,13 +217,9 @@ def interfaces():
# legacy single-value behavior (keeps the last IPv4 seen)
ip_address = addr.address
netmask = addr.netmask
ip_addresses.append(
{"family": "ipv4", "address": addr.address, "netmask": addr.netmask or None}
)
ip_addresses.append({"family": "ipv4", "address": addr.address, "netmask": addr.netmask or None})
elif addr.family == socket.AF_INET6:
ip_addresses.append(
{"family": "ipv6", "address": addr.address, "netmask": addr.netmask or None}
)
ip_addresses.append({"family": "ipv6", "address": addr.address, "netmask": addr.netmask or None})
if addr.family == psutil.AF_LINK:
mac_address = addr.address
if interface.startswith("tap"):
@@ -283,7 +278,7 @@ def interfaces():
"vmware",
"virtualbox",
"gns3",
"veb"
"veb",
):
if result["name"].lower().startswith(special_interface):
result["special"] = True
+11 -31
View File
@@ -11,6 +11,7 @@ log = logging.getLogger(__name__)
class FilterValidationError(Exception):
"""Raised when packet filter parameters fail validation."""
pass
@@ -46,7 +47,7 @@ def validate_bpf_syntax(bpf_expression: str) -> Dict[str, Optional[str]]:
# Strip "tcpdump: " prefix
for prefix in ["tcpdump: "]:
if line.startswith(prefix):
line = line[len(prefix):]
line = line[len(prefix) :]
error_lines.append(line)
error_msg = " ".join(error_lines) if error_lines else "Invalid BPF expression"
log.warning("BPF syntax validation failed: %s", error_msg)
@@ -56,10 +57,7 @@ def validate_bpf_syntax(bpf_expression: str) -> Dict[str, Optional[str]]:
return {"valid": True, "error": None}
except FileNotFoundError:
log.warning(
"tcpdump not found, skipping BPF syntax validation. "
"Install tcpdump to enable BPF validation."
)
log.warning("tcpdump not found, skipping BPF syntax validation. Install tcpdump to enable BPF validation.")
return {"valid": True, "error": None}
except Exception as e:
@@ -85,31 +83,17 @@ def validate_filter_parameters(filter_type: str, values: List[Any]) -> None:
"params_count": 1,
"ranges": [(-1, 32767)], # min, max
"names": ["Frequency"],
"units": ["th packet"]
},
"packet_loss": {
"params_count": 1,
"ranges": [(0, 100)],
"names": ["Chance"],
"units": ["%"]
"units": ["th packet"],
},
"packet_loss": {"params_count": 1, "ranges": [(0, 100)], "names": ["Chance"], "units": ["%"]},
"delay": {
"params_count": 2, # latency, jitter
"ranges": [(1, 32767), (0, 32767)], # ubridge rejects latency <= 0
"names": ["Latency", "Jitter"],
"units": ["ms", "ms"]
"units": ["ms", "ms"],
},
"corrupt": {
"params_count": 1,
"ranges": [(0, 100)],
"names": ["Chance"],
"units": ["%"]
},
"bpf": {
"params_count": 1,
"is_text": True,
"names": ["Filters"]
}
"corrupt": {"params_count": 1, "ranges": [(0, 100)], "names": ["Chance"], "units": ["%"]},
"bpf": {"params_count": 1, "is_text": True, "names": ["Filters"]},
}
if filter_type not in VALIDATION_RULES:
@@ -119,18 +103,14 @@ def validate_filter_parameters(filter_type: str, values: List[Any]) -> None:
# Check parameter count
if len(values) != rules["params_count"]:
raise FilterValidationError(
f"{filter_type} expects {rules['params_count']} parameter(s), got {len(values)}"
)
raise FilterValidationError(f"{filter_type} expects {rules['params_count']} parameter(s), got {len(values)}")
# Validate each parameter
for i, value in enumerate(values):
if rules.get("is_text"):
# Text validation (BPF)
if not isinstance(value, str):
raise FilterValidationError(
f"{filter_type} parameter {rules['names'][i]} must be a string"
)
raise FilterValidationError(f"{filter_type} parameter {rules['names'][i]} must be a string")
# Validate BPF syntax using tshark (same method as gns3_copilot)
# The value may be a multi-line string; each line becomes a
@@ -251,4 +231,4 @@ def validate_all_filters(filters: Dict[str, List[Any]]) -> None:
if not values or (isinstance(values, list) and len(values) == 0):
continue
validate_filter_parameters(filter_type, values)
validate_filter_parameters(filter_type, values)
+1 -1
View File
@@ -73,7 +73,7 @@ def extract_client_info(scope: Dict[str, Any], auth_service_instance: Optional[A
"path": path,
"method": method,
"username": username,
"user_info": user_info
"user_info": user_info,
}
+4 -11
View File
@@ -15,7 +15,7 @@ import argparse
def _find_base_dir():
"""Find the project root directory."""
if hasattr(sys, '_MEIPASS'):
if hasattr(sys, "_MEIPASS"):
return os.path.dirname(sys.executable)
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -71,9 +71,7 @@ def uninstall(packages, yes=False):
for package in packages:
try:
result = subprocess.run(
[sys.executable, "-m", "pip", "uninstall", "-y", package],
capture_output=True,
text=True
[sys.executable, "-m", "pip", "uninstall", "-y", package], capture_output=True, text=True
)
if result.returncode == 0:
print(f" Removed: {package}")
@@ -87,13 +85,8 @@ def uninstall(packages, yes=False):
def main():
parser = argparse.ArgumentParser(
description="Uninstall AI Features dependencies (AI Copilot + MCP)"
)
parser.add_argument(
"-y", "--yes", action="store_true",
help="Automatically confirm uninstallation"
)
parser = argparse.ArgumentParser(description="Uninstall AI Features dependencies (AI Copilot + MCP)")
parser.add_argument("-y", "--yes", action="store_true", help="Automatically confirm uninstallation")
args = parser.parse_args()
packages = get_ai_packages()
+1 -4
View File
@@ -7,10 +7,7 @@ import argparse
import re
UUID_PATTERN = re.compile(
r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$',
re.IGNORECASE
)
UUID_PATTERN = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE)
def validate_uuid(value: str) -> str:
+13 -35
View File
@@ -41,7 +41,7 @@ async def websocket_proxy(
target_url: str,
requested_protocols: list = None,
buffer_size: int = 65536,
timeout: Optional[float] = None
timeout: Optional[float] = None,
) -> None:
"""
Proxy binary WebSocket data between client and target WebSocket server.
@@ -116,28 +116,20 @@ async def websocket_proxy(
if sys.version_info >= (3, 11, 0):
aws = [
asyncio.create_task(forward_client_to_target(target_ws)),
asyncio.create_task(forward_target_to_client(target_ws))
asyncio.create_task(forward_target_to_client(target_ws)),
]
else:
aws = [
forward_client_to_target(target_ws),
forward_target_to_client(target_ws)
]
aws = [forward_client_to_target(target_ws), forward_target_to_client(target_ws)]
try:
done, pending = await asyncio.wait(
aws,
return_when=asyncio.FIRST_COMPLETED
)
done, pending = await asyncio.wait(aws, return_when=asyncio.FIRST_COMPLETED)
except Exception as e:
log.error(f"asyncio.wait raised exception: {e}")
# Check for exceptions
for task in done:
if task.exception():
log.warning(
f"WebSocket proxy task exception: {task.exception()}"
)
log.warning(f"WebSocket proxy task exception: {task.exception()}")
# Cancel pending tasks
for task in pending:
@@ -145,16 +137,10 @@ async def websocket_proxy(
except aiohttp.ClientError as e:
log.error(f"WebSocket proxy connection error: {e}")
await client_ws.close(
code=status.WS_1011_INTERNAL_ERROR,
reason=f"Proxy connection failed: {e}"
)
await client_ws.close(code=status.WS_1011_INTERNAL_ERROR, reason=f"Proxy connection failed: {e}")
except Exception as e:
log.error(f"WebSocket proxy unexpected error: {e}")
await client_ws.close(
code=status.WS_1011_INTERNAL_ERROR,
reason=str(e)
)
await client_ws.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(e))
async def websocket_proxy_with_manual_accept(
@@ -162,7 +148,7 @@ async def websocket_proxy_with_manual_accept(
target_url: str,
requested_protocols: list = None,
buffer_size: int = 65536,
timeout: Optional[float] = None
timeout: Optional[float] = None,
) -> None:
"""
Proxy binary WebSocket data between client and target WebSocket server.
@@ -181,7 +167,7 @@ async def websocket_proxy_with_manual_accept(
Raises:
aiohttp.ClientError: If connection to target fails
"""
client_info = f"{client_ws.client.host}:{client_ws.client.port}" if hasattr(client_ws, 'client') else "unknown"
client_info = f"{client_ws.client.host}:{client_ws.client.port}" if hasattr(client_ws, "client") else "unknown"
async def forward_client_to_target(target_ws):
"""Client → Target: Forward binary WebSocket data."""
@@ -252,20 +238,14 @@ async def websocket_proxy_with_manual_accept(
if sys.version_info >= (3, 11, 0):
aws = [
asyncio.create_task(forward_client_to_target(target_ws)),
asyncio.create_task(forward_target_to_client(target_ws))
asyncio.create_task(forward_target_to_client(target_ws)),
]
else:
aws = [
forward_client_to_target(target_ws),
forward_target_to_client(target_ws)
]
aws = [forward_client_to_target(target_ws), forward_target_to_client(target_ws)]
log.info(f"About to call asyncio.wait with {len(aws)} tasks")
try:
done, pending = await asyncio.wait(
aws,
return_when=asyncio.FIRST_COMPLETED
)
done, pending = await asyncio.wait(aws, return_when=asyncio.FIRST_COMPLETED)
log.info(f"asyncio.wait returned. done={len(done)}, pending={len(pending)}")
except Exception as e:
log.error(f"asyncio.wait raised exception: {e}")
@@ -274,9 +254,7 @@ async def websocket_proxy_with_manual_accept(
# Check for exceptions
for task in done:
if task.exception():
log.warning(
f"WebSocket proxy task exception: {task.exception()}"
)
log.warning(f"WebSocket proxy task exception: {task.exception()}")
# Cancel pending tasks
for task in pending:
@@ -1,4 +1,3 @@
# NOTE: this patches the standard zipfile module
from . import _zipfile
@@ -7,4 +6,3 @@ from zipfile import (
ZIP_ZSTANDARD,
ZSTANDARD_VERSION,
)
@@ -2,7 +2,6 @@ import functools
class patch:
originals = {}
def __init__(self, host, name):
@@ -17,4 +16,3 @@ class patch:
setattr(self.host, self.name, func)
return func
+15 -13
View File
@@ -6,16 +6,16 @@ from ._patcher import patch
zipfile.ZIP_ZSTANDARD = 93
zipfile.compressor_names[zipfile.ZIP_ZSTANDARD] = 'zstandard'
zipfile.compressor_names[zipfile.ZIP_ZSTANDARD] = "zstandard"
zipfile.ZSTANDARD_VERSION = 20
@patch(zipfile, '_check_compression')
@patch(zipfile, "_check_compression")
def zstd_check_compression(compression):
if compression == zipfile.ZIP_ZSTANDARD:
pass
else:
patch.originals['_check_compression'](compression)
patch.originals["_check_compression"](compression)
class ZstdDecompressObjWrapper:
@@ -23,42 +23,44 @@ class ZstdDecompressObjWrapper:
self.o = o
def __getattr__(self, attr):
if attr == 'eof':
if attr == "eof":
return False
return getattr(self.o, attr)
@patch(zipfile, '_get_decompressor')
@patch(zipfile, "_get_decompressor")
def zstd_get_decompressor(compress_type):
if compress_type == zipfile.ZIP_ZSTANDARD:
return ZstdDecompressObjWrapper(zstd.ZstdDecompressor(max_window_size=2147483648).decompressobj())
else:
return patch.originals['_get_decompressor'](compress_type)
return patch.originals["_get_decompressor"](compress_type)
if 'compresslevel' in inspect.signature(zipfile._get_compressor).parameters:
@patch(zipfile, '_get_compressor')
if "compresslevel" in inspect.signature(zipfile._get_compressor).parameters:
@patch(zipfile, "_get_compressor")
def zstd_get_compressor(compress_type, compresslevel=None):
if compress_type == zipfile.ZIP_ZSTANDARD:
if compresslevel is None:
compresslevel = 3
return zstd.ZstdCompressor(level=compresslevel, threads=-1).compressobj()
else:
return patch.originals['_get_compressor'](compress_type, compresslevel=compresslevel)
return patch.originals["_get_compressor"](compress_type, compresslevel=compresslevel)
else:
@patch(zipfile, '_get_compressor')
@patch(zipfile, "_get_compressor")
def zstd_get_compressor(compress_type, compresslevel=None):
if compress_type == zipfile.ZIP_ZSTANDARD:
if compresslevel is None:
compresslevel = 3
return zstd.ZstdCompressor(level=compresslevel, threads=-1).compressobj()
else:
return patch.originals['_get_compressor'](compress_type)
return patch.originals["_get_compressor"](compress_type)
@patch(zipfile.ZipInfo, 'FileHeader')
@patch(zipfile.ZipInfo, "FileHeader")
def zstd_FileHeader(self, zip64=None):
if self.compress_type == zipfile.ZIP_ZSTANDARD:
self.create_version = max(self.create_version, zipfile.ZSTANDARD_VERSION)
self.extract_version = max(self.extract_version, zipfile.ZSTANDARD_VERSION)
return patch.originals['FileHeader'](self, zip64=zip64)
return patch.originals["FileHeader"](self, zip64=zip64)