324 lines
7.5 KiB
Python
324 lines
7.5 KiB
Python
import asyncio
|
|
import json
|
|
import shutil
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
WORKSPACE_ROOT = Path(tempfile.gettempdir()) / "theplug-lsp"
|
|
CONFIG_FILE = "pyrightconfig.json"
|
|
ENTRY_FILE = "proxy_addon.py"
|
|
|
|
STUB_FILES = {
|
|
"typings/mitmproxy/__init__.pyi": """from . import (
|
|
command,
|
|
ctx,
|
|
dns,
|
|
http,
|
|
tcp,
|
|
tls,
|
|
udp,
|
|
websocket,
|
|
)
|
|
|
|
__all__ = [
|
|
"command",
|
|
"ctx",
|
|
"dns",
|
|
"http",
|
|
"tcp",
|
|
"tls",
|
|
"udp",
|
|
"websocket",
|
|
]
|
|
""",
|
|
"typings/mitmproxy/ctx.pyi": """class _Logger:
|
|
def debug(self, message: object) -> None: ...
|
|
def info(self, message: object) -> None: ...
|
|
def warn(self, message: object) -> None: ...
|
|
def error(self, message: object) -> None: ...
|
|
def alert(self, message: object) -> None: ...
|
|
|
|
|
|
class _Options:
|
|
def __getattr__(self, name: str) -> object: ...
|
|
|
|
|
|
log: _Logger
|
|
options: _Options
|
|
""",
|
|
"typings/mitmproxy/http.pyi": """from typing import Any
|
|
|
|
|
|
class Headers(dict[str, str]):
|
|
...
|
|
|
|
|
|
class Request:
|
|
host: str
|
|
path: str
|
|
headers: Headers
|
|
query: dict[str, str]
|
|
content: bytes
|
|
|
|
|
|
class Response:
|
|
headers: Headers
|
|
status_code: int
|
|
content: bytes
|
|
|
|
@classmethod
|
|
def make(
|
|
cls,
|
|
status_code: int,
|
|
content: bytes | str = b"",
|
|
headers: dict[str, str] | None = None,
|
|
) -> "Response": ...
|
|
|
|
|
|
class WebSocketMessage:
|
|
content: bytes
|
|
from_client: bool
|
|
|
|
|
|
class WebSocketData:
|
|
messages: list[WebSocketMessage]
|
|
|
|
|
|
class HTTPFlow:
|
|
request: Request
|
|
response: Response | None
|
|
websocket: WebSocketData | None
|
|
metadata: dict[str, Any]
|
|
""",
|
|
"typings/mitmproxy/tcp.pyi": """class TCPMessage:
|
|
content: bytes
|
|
from_client: bool
|
|
|
|
|
|
class TCPFlow:
|
|
messages: list[TCPMessage]
|
|
""",
|
|
"typings/mitmproxy/dns.pyi": """class Question:
|
|
name: str
|
|
type: int
|
|
|
|
|
|
class ResourceRecord:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
name: str,
|
|
type: int,
|
|
class_: int,
|
|
ttl: int,
|
|
data: bytes,
|
|
) -> None: ...
|
|
|
|
|
|
class Message:
|
|
question: Question | None
|
|
|
|
def succeed(self, records: list[ResourceRecord]) -> "Message": ...
|
|
|
|
|
|
class _Types:
|
|
A: int
|
|
|
|
|
|
class _Classes:
|
|
IN: int
|
|
|
|
|
|
types: _Types
|
|
classes: _Classes
|
|
|
|
|
|
class DNSFlow:
|
|
request: Message
|
|
response: Message | None
|
|
""",
|
|
"typings/mitmproxy/command.pyi": """from collections.abc import Callable
|
|
from typing import TypeVar
|
|
|
|
F = TypeVar("F", bound=Callable[..., object])
|
|
|
|
|
|
def command(name: str) -> Callable[[F], F]: ...
|
|
""",
|
|
"typings/mitmproxy/websocket.pyi": """class WebSocketMessage:
|
|
content: bytes
|
|
from_client: bool
|
|
""",
|
|
"typings/mitmproxy/tls.pyi": """class ClientHelloData: ...
|
|
""",
|
|
"typings/mitmproxy/udp.pyi": """class UDPMessage:
|
|
content: bytes
|
|
from_client: bool
|
|
""",
|
|
}
|
|
|
|
|
|
class LspFrameBuffer:
|
|
def __init__(self) -> None:
|
|
self.buffer = bytearray()
|
|
|
|
def feed(self, chunk: bytes) -> list[str]:
|
|
self.buffer.extend(chunk)
|
|
messages: list[str] = []
|
|
while True:
|
|
header_end = self.buffer.find(b"\r\n\r\n")
|
|
if header_end < 0:
|
|
return messages
|
|
|
|
header = self.buffer[:header_end].decode("ascii", errors="ignore")
|
|
size = content_length(header)
|
|
if size is None:
|
|
del self.buffer[: header_end + 4]
|
|
continue
|
|
|
|
frame_end = header_end + 4 + size
|
|
if len(self.buffer) < frame_end:
|
|
return messages
|
|
|
|
body = bytes(self.buffer[header_end + 4 : frame_end])
|
|
del self.buffer[:frame_end]
|
|
messages.append(body.decode("utf-8"))
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class PythonEditorSession:
|
|
root_path: Path
|
|
file_path: Path
|
|
|
|
@property
|
|
def root_uri(self) -> str:
|
|
return self.root_path.as_uri()
|
|
|
|
@property
|
|
def file_uri(self) -> str:
|
|
return self.file_path.as_uri()
|
|
|
|
|
|
def content_length(header: str) -> int | None:
|
|
for line in header.split("\r\n"):
|
|
if ":" not in line:
|
|
continue
|
|
name, value = line.split(":", 1)
|
|
if name.lower() == "content-length":
|
|
return int(value.strip())
|
|
return None
|
|
|
|
|
|
def encode_message(message: str) -> bytes:
|
|
payload = message.encode("utf-8")
|
|
header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
|
|
return header + payload
|
|
|
|
|
|
def lsp_binary() -> str | None:
|
|
return shutil.which("basedpyright-langserver")
|
|
|
|
|
|
def get_session() -> PythonEditorSession:
|
|
root_path = WORKSPACE_ROOT
|
|
file_path = root_path / ENTRY_FILE
|
|
|
|
root_path.mkdir(parents=True, exist_ok=True)
|
|
file_path.touch(exist_ok=True)
|
|
|
|
config = {
|
|
"include": [ENTRY_FILE],
|
|
"stubPath": "typings",
|
|
"typeCheckingMode": "basic",
|
|
"pythonVersion": "3.14",
|
|
"reportMissingTypeStubs": False,
|
|
"reportMissingImports": "warning",
|
|
"useLibraryCodeForTypes": True,
|
|
}
|
|
(root_path / CONFIG_FILE).write_text(
|
|
json.dumps(config, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
for relative_path, content in STUB_FILES.items():
|
|
path = root_path / relative_path
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(content, encoding="utf-8")
|
|
|
|
return PythonEditorSession(root_path=root_path, file_path=file_path)
|
|
|
|
|
|
async def open_process(session: PythonEditorSession) -> asyncio.subprocess.Process:
|
|
binary = lsp_binary()
|
|
if binary is None:
|
|
raise RuntimeError("basedpyright-langserver is not installed")
|
|
|
|
return await asyncio.create_subprocess_exec(
|
|
binary,
|
|
"--stdio",
|
|
cwd=session.root_path,
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
|
|
|
|
async def forward_client(websocket, writer: asyncio.StreamWriter) -> None:
|
|
while True:
|
|
message = await websocket.receive_text()
|
|
writer.write(encode_message(message))
|
|
await writer.drain()
|
|
|
|
|
|
async def forward_server(reader: asyncio.StreamReader, websocket) -> None:
|
|
frames = LspFrameBuffer()
|
|
while True:
|
|
chunk = await reader.read(65536)
|
|
if not chunk:
|
|
return
|
|
for message in frames.feed(chunk):
|
|
await websocket.send_text(message)
|
|
|
|
|
|
async def drain_errors(reader: asyncio.StreamReader) -> None:
|
|
while True:
|
|
chunk = await reader.read(4096)
|
|
if not chunk:
|
|
return
|
|
|
|
|
|
async def bridge(websocket, session: PythonEditorSession) -> None:
|
|
process = await open_process(session)
|
|
stdin = process.stdin
|
|
stdout = process.stdout
|
|
stderr = process.stderr
|
|
|
|
if stdin is None or stdout is None or stderr is None:
|
|
process.kill()
|
|
await process.wait()
|
|
raise RuntimeError("basedpyright-langserver did not expose stdio pipes")
|
|
|
|
send_task = asyncio.create_task(forward_client(websocket, stdin))
|
|
receive_task = asyncio.create_task(forward_server(stdout, websocket))
|
|
error_task = asyncio.create_task(drain_errors(stderr))
|
|
|
|
try:
|
|
done, pending = await asyncio.wait(
|
|
{send_task, receive_task, error_task},
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
for task in done:
|
|
task.result()
|
|
finally:
|
|
if process.returncode is None:
|
|
stdin.close()
|
|
await stdin.wait_closed()
|
|
try:
|
|
await asyncio.wait_for(process.wait(), timeout=1)
|
|
except TimeoutError:
|
|
process.kill()
|
|
await process.wait()
|