diff --git a/rcon/console.py b/rcon/console.py index 5cca72f..faf21e8 100644 --- a/rcon/console.py +++ b/rcon/console.py @@ -140,6 +140,7 @@ def rconcmd( port: int, passwd: str, *, + timeout: float | None = None, prompt: str = PROMPT ): """Initialize the console.""" @@ -152,7 +153,7 @@ def rconcmd( prompt = prompt.format(host=host, port=port) - with client_cls(host, port) as client: + with client_cls(host, port, timeout=timeout) as client: try: passwd = login(client, passwd) except EOFError: diff --git a/rcon/rconshell.py b/rcon/rconshell.py index e13b3b7..c5d5d6e 100644 --- a/rcon/rconshell.py +++ b/rcon/rconshell.py @@ -34,6 +34,10 @@ def get_args() -> Namespace: '-p', '--prompt', default=PROMPT, metavar='PS1', help='the shell prompt' ) + parser.add_argument( + '-t', '--timeout', type=float, metavar='seconds', + help='connection timeout in seconds' + ) return parser.parse_args() @@ -50,7 +54,14 @@ def run() -> None: host = port = passwd = None with CommandHistory(LOGGER): - rconcmd(client_cls, host, port, passwd, prompt=args.prompt) + rconcmd( + client_cls, + host, + port, + passwd, + timeout=args.timeout, + prompt=args.prompt + ) def main() -> int: