From 10a41118b678d897bc06caf18cd35eecdc7f6865 Mon Sep 17 00:00:00 2001
From: themylogin <themylogin@gmail.com>
Date: Mon, 9 Dec 2024 13:56:08 +0100
Subject: [PATCH] `private_methods`

---
 truenas_api_client/__init__.py | 19 +++++++++++++------
 truenas_api_client/legacy.py   |  4 ++--
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/truenas_api_client/__init__.py b/truenas_api_client/__init__.py
index 57ae351..5784a3a 100644
--- a/truenas_api_client/__init__.py
+++ b/truenas_api_client/__init__.py
@@ -65,8 +65,8 @@
 class Client:
     """Implicit wrapper of either a `JSONRPCClient` or a `LegacyClient`."""
 
-    def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False,
-                 call_timeout: float | UndefinedType=undefined, verify_ssl=True):
+    def __init__(self, uri: str | None=None, reserved_ports=False, private_methods=False, py_exceptions=False,
+                 log_py_exceptions=False, call_timeout: float | UndefinedType=undefined, verify_ssl=True):
         """Initialize either a `JSONRPCClient` or a `LegacyClient`.
 
         Use `JSONRPCClient` unless `uri` ends with '/websocket'.
@@ -74,6 +74,7 @@ def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=Fal
         Args:
             uri: The address to connect to. Defaults to the local middlewared socket.
             reserved_ports: `True` if the local socket should use a reserved port.
+            private_methods: `True` if calling private methods should be allowed
             py_exceptions: `True` if the server should include exception objects in
                 `message['error']['data']['py_exception']`.
             log_py_exceptions: `True` if exception tracebacks from API calls should be logged.
@@ -90,7 +91,8 @@ def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=Fal
         else:
             client_class = JSONRPCClient
 
-        self.__client = client_class(uri, reserved_ports, py_exceptions, log_py_exceptions, call_timeout, verify_ssl)
+        self.__client = client_class(uri, reserved_ports, private_methods, py_exceptions, log_py_exceptions,
+                                     call_timeout, verify_ssl)
 
     def __getattr__(self, item):
         return getattr(self.__client, item)
@@ -397,13 +399,14 @@ class JSONRPCClient:
     Keeps track of the calls made, jobs submitted, and callbacks. Maintains a websocket connection using a `WSClient`.
 
     """
-    def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False,
-                 call_timeout: float | UndefinedType=undefined, verify_ssl=True):
+    def __init__(self, uri: str | None = None, reserved_ports=False, private_methods=False, py_exceptions=False,
+                 log_py_exceptions=False, call_timeout: float | UndefinedType = undefined, verify_ssl=True):
         """Initialize a `JSONRPCClient`.
 
         Args:
             uri: The address to connect to. Defaults to the local middlewared socket.
             reserved_ports: `True` if the local socket should use a reserved port.
+            private_methods: `True` if calling private methods should be allowed
             py_exceptions: `True` if the server should include exception objects in
                 `message['error']['data']['py_exception']`.
             log_py_exceptions: `True` if exception tracebacks from API calls should be logged.
@@ -425,6 +428,7 @@ def __init__(self, uri: str | None=None, reserved_ports=False, py_exceptions=Fal
         self._jobs: defaultdict[str, _JobDict] = defaultdict(dict)  # type: ignore
         self._jobs_lock = Lock()
         self._jobs_watching = False
+        self._private_methods = private_methods
         self._py_exceptions = py_exceptions
         self._log_py_exceptions = log_py_exceptions
         self._call_timeout = call_timeout
@@ -583,7 +587,10 @@ def _run_callback(self, event: _Payload, args: Iterable[str], kwargs: Collection
 
     def on_open(self):
         """Make an API call to `core.set_options` to configure how middlewared sends its responses."""
-        self._set_options_call = self.call("core.set_options", {"py_exceptions": self._py_exceptions}, background=True)
+        self._set_options_call = self.call("core.set_options", {
+            "private_methods": self._private_methods,
+            "py_exceptions": self._py_exceptions,
+        }, background=True)
 
     def on_close(self, code: int, reason: str | None=None):
         """Close this `JSONRPCClient` in response to the `WebSocketApp` closing.
diff --git a/truenas_api_client/legacy.py b/truenas_api_client/legacy.py
index 5a42a8b..7298d47 100644
--- a/truenas_api_client/legacy.py
+++ b/truenas_api_client/legacy.py
@@ -192,8 +192,8 @@ def result(self):
 
 
 class LegacyClient:
-    def __init__(self, uri=None, reserved_ports=False, py_exceptions=False, log_py_exceptions=False,
-                 call_timeout: float | UndefinedType=undefined, verify_ssl=True):
+    def __init__(self, uri=None, reserved_ports=False, private_methods=False, py_exceptions=False,
+                 log_py_exceptions=False, call_timeout: float | UndefinedType=undefined, verify_ssl=True):
         """
         Arguments:
            :reserved_ports(bool): should the local socket used a reserved port