From 7e78a784ba320fb08834944fee4913703f2d12c6 Mon Sep 17 00:00:00 2001
From: Richard Neumann <mail@richard-neumann.de>
Date: Fri, 13 Jan 2023 23:59:34 +0100
Subject: [PATCH] Add optional timeout to RCON shell

---
 rcon/console.py   |  3 ++-
 rcon/rconshell.py | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

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: