Add helpers to set TCP parameters or enable/disable keepalives

Signed-off-by: Olivier Gayot <olivier.gayot@sigexec.com>

1 files changed, 27 insertions, 27 deletions

diff --git a/solter.py b/solter.py
index c224ae6..355b7dd 100644
--- a/solter.py
+++ b/solter.py
@@ -14,49 +14,49 @@ TCP_KEEPINTVL = 5
 TCP_KEEPCNT = 6
 
 
-def get_commands_enable_keepalive(sockfd: int, **kwargs):
+def _get_commands_enable_disable_keepalive(sockfd: int, enable: bool, **kwargs):
+    """ Return GDB instructions to enable or disable keepalive on a given socket """
     return (
             "set $rsp -= sizeof(int)",
-            "set {int}$rsp = 1",
+            f"set {{int}}$rsp = {1 if enable else 0}",
             f"call (int)setsockopt({sockfd}, {SO_SOCKET}, {SO_KEEPALIVE}, $rsp, sizeof(int))",
             "set $rsp += sizeof(int)",
     )
 
 
-def get_commands_disable_keepalive(sockfd: int, **kwargs):
-    return (
-            "set $rsp -= sizeof(int)",
-            "set {int}$rsp = 0",
-            f"call (int)setsockopt({sockfd}, {SO_SOCKET}, {SO_KEEPALIVE}, $rsp, sizeof(int))",
-            "set $rsp += sizeof(int)",
-    )
-
-
-def get_commands_change_keepalive_time(sockfd: int, value: int):
+def _get_commands_change_tcp_option(sockfd: int, value: int, param: int, **kwargs):
+    """ Return GDB instructions to change a given TCP option on a given socket """
     return (
             "set $rsp -= sizeof(int)",
             f"set {{int}}$rsp = {value}",
-            f"call (int)setsockopt({sockfd}, {SO_TCP}, {TCP_KEEPIDLE}, $rsp, sizeof(int))",
+            f"call (int)setsockopt({sockfd}, {SO_TCP}, {param}, $rsp, sizeof(int))",
             "set $rsp += sizeof(int)",
             )
 
 
-def get_commands_change_keepalive_intvl(sockfd: int, value: int):
-    return (
-            "set $rsp -= sizeof(int)",
-            f"set {{int}}$rsp = {value}",
-            f"call (int)setsockopt({sockfd}, {SO_TCP}, {TCP_KEEPINTVL}, $rsp, sizeof(int))",
-            "set $rsp += sizeof(int)",
-            )
+def get_commands_enable_keepalive(**kwargs):
+    """ Return GDB commands to enable keepalive on a given socket """
+    return _get_commands_enable_disable_keepalive(enable=True, **kwargs)
 
 
-def get_commands_change_keepalive_count(sockfd: int, value: int):
-    return (
-            "set $rsp -= sizeof(int)",
-            f"set {{int}}$rsp = {value}",
-            f"call (int)setsockopt({sockfd}, {SO_TCP}, {TCP_KEEPCNT}, $rsp, sizeof(int))",
-            "set $rsp += sizeof(int)",
-            )
+def get_commands_disable_keepalive(**kwargs):
+    """ Return GDB commands to disable keepalive on a given socket """
+    return _get_commands_enable_disable_keepalive(enable=False, **kwargs)
+
+
+def get_commands_change_keepalive_time(**kwargs):
+    """ Return GDB commands to set the TCP keepalive time on a given socket """
+    return _get_commands_change_tcp_option(param=TCP_KEEPIDLE, **kwargs)
+
+
+def get_commands_change_keepalive_intvl(**kwargs):
+    """ Return GDB commands to set the TCP keepalive interval on a given socket """
+    return _get_commands_change_tcp_option(param=TCP_KEEPINTVL, **kwargs)
+
+
+def get_commands_change_keepalive_count(**kwargs):
+    """ Return GDB commands to set the TCP keepalive count on a given socket """
+    return _get_commands_change_tcp_option(param=TCP_KEEPCNT, **kwargs)
 
 
 def main(args):