Skip to content

Commit

Permalink
♻️ ground work for forcibly closing connections
Browse files Browse the repository at this point in the history
haliphax committed Apr 10, 2024
1 parent e27e6d5 commit 8c55788
Showing 4 changed files with 14 additions and 6 deletions.
7 changes: 4 additions & 3 deletions xthulu/ssh/__init__.py
Original file line number Diff line number Diff line change
@@ -8,7 +8,7 @@
from tracemalloc import start

# 3rd party
from asyncssh import listen
from asyncssh import SSHAcceptor, listen

# local
from ..configuration import get_config
@@ -20,7 +20,7 @@
from .server import SSHServer


async def start_server():
async def start_server() -> SSHAcceptor:
"""Start the SSH server."""

register_encodings()
@@ -49,4 +49,5 @@ async def start_server():
start()

await res.db.set_bind(res.db.bind)
await listen(**kwargs)

return await listen(**kwargs)
4 changes: 2 additions & 2 deletions xthulu/ssh/context/__init__.py
Original file line number Diff line number Diff line change
@@ -21,7 +21,7 @@
from ...models import User
from ...scripting import load_userland_module
from ..console import XthuluConsole
from ..exceptions import Goto, ProcessClosing
from ..exceptions import Goto, ProcessClosing, ProcessForciblyClosed
from ..structs import Script
from .lock_manager import _LockManager
from .logger_adapter import ContextLoggerAdapter
@@ -275,7 +275,7 @@ async def runscript(self, script: Script) -> Any:
mod = load_userland_module(script.name)
main: Callable[..., Any] = getattr(mod, "main")
return await main(self, *script.args, **script.kwargs)
except (ProcessClosing, Goto):
except (ProcessClosing, ProcessForciblyClosed, Goto):
raise
except Exception:
message = f"Exception in script {script.name}"
4 changes: 4 additions & 0 deletions xthulu/ssh/exceptions.py
Original file line number Diff line number Diff line change
@@ -20,3 +20,7 @@ def __init__(self, script: str, *args, **kwargs):

class ProcessClosing(Exception):
"""Thrown when the `asyncssh.SSHServerProcess` is closing"""


class ProcessForciblyClosed(Exception):
"""Thrown when the process is being forcibly closed by the server"""
5 changes: 4 additions & 1 deletion xthulu/ssh/process_factory.py
Original file line number Diff line number Diff line change
@@ -12,7 +12,7 @@
from ..events.structs import EventData
from .console import XthuluConsole
from .context import SSHContext
from .exceptions import Goto, ProcessClosing
from .exceptions import Goto, ProcessClosing, ProcessForciblyClosed
from .structs import Script


@@ -68,6 +68,9 @@ async def input_loop():
# process is likely closing
break

except ProcessForciblyClosed:
break

except TimeoutError:
cx.log.warn("Timed out")
cx.echo("\n\n[bright_white on red] TIMED OUT [/]\n\n")

0 comments on commit 8c55788

Please sign in to comment.