rpcclient: add Allocated

src/rpcclient/rpcclient/allocated.py

@@ -0,0 +1,26 @@
+from abc import abstractmethod
+
+
+class Allocated:
+    """ resource allocated on remote host that needs to be free """
+
+    def __init__(self):
+        self._deallocated = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.deallocate()
+
+    def __del__(self):
+        self.deallocate()
+
+    @abstractmethod
+    def _deallocate(self):
+        pass
+
+    def deallocate(self):
+        if not self._deallocated:
+            self._deallocated = True
+            self._deallocate()

src/rpcclient/rpcclient/darwin/fs.py

@@ -29,6 +29,7 @@ def __iter__(self):
             self._client.errno = 0
             direntp = self._client.symbols.readdir(self._dirp)
             if not direntp:
+                self.deallocate()
                 if self._client.errno:
                     raise BadReturnValueError(f'Failed scanning dir: {self._client.last_error}')
                 break
@@ -37,7 +38,7 @@ def __iter__(self):
                 continue
             yield DarwinDirEntry(self.path, entry, self._client)
 
-    def closedir(self):
+    def _deallocate(self):
         self._client.symbols.closedir(self._dirp)
 
 

src/rpcclient/rpcclient/darwin/network.py

@@ -4,6 +4,7 @@
 
 from rpcclient.exceptions import BadReturnValueError, RpcClientException
 from rpcclient.network import Network
+from rpcclient.allocated import Allocated
 
 logger = logging.getLogger(__name__)
 
@@ -46,19 +47,14 @@ def __repr__(self):
         return result
 
 
-class WifiInterface:
+class WifiInterface(Allocated):
     def __init__(self, client, interface, wifi_manager_client, device):
+        super().__init__()
         self._client = client
         self._interface = interface
         self._wifi_manager_client = wifi_manager_client
         self._device = device
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
     def scan(self) -> List[WifiNetwork]:
         """ perform wifi scan """
         result = []
@@ -96,7 +92,7 @@ def turn_off(self):
     def is_on(self) -> bool:
         return self._client.symbols.WiFiDeviceClientGetPower(self._device) == 1
 
-    def close(self):
+    def _deallocate(self):
         self._client.symbols.Apple80211Close(self._interface)
 
 

src/rpcclient/rpcclient/darwin/scpreferences.py

@@ -1,22 +1,15 @@
 import typing
 
 from rpcclient.exceptions import RpcClientException
+from rpcclient.allocated import Allocated
 
 
-class SCPreferencesObject:
+class SCPreferencesObject(Allocated):
     def __init__(self, client, ref):
+        super().__init__()
         self._client = client
         self._ref = ref
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        try:
-            yield
-        finally:
-            self.release()
-
     @property
     def keys(self) -> typing.List[str]:
         return self._client.symbols.SCPreferencesCopyKeyList(self._ref).py
@@ -40,7 +33,7 @@ def to_dict(self) -> typing.Mapping:
             result[k] = self.get(k)
         return result
 
-    def release(self):
+    def _deallocate(self):
         self._client.CFRelase(self._ref)
 
     def _commit(self):

src/rpcclient/rpcclient/fs.py

@@ -3,6 +3,7 @@
 
 from rpcclient.common import path_to_str
 from rpcclient.exceptions import InvalidArgumentError, BadReturnValueError
+from rpcclient.allocated import Allocated
 from rpcclient.structs.consts import O_RDONLY, O_WRONLY, O_CREAT, O_TRUNC, S_IFMT, S_IFDIR, O_RDWR, SEEK_CUR, S_IFREG, \
     DT_LNK, DT_UNKNOWN, S_IFLNK, DT_REG, DT_DIR
 
@@ -76,59 +77,30 @@ def _fetch_stat(self, follow_symlinks):
         raise NotImplementedError()
 
 
-class ScandirIterator:
+class ScandirIterator(Allocated):
     def __init__(self, path, dirp, client):
+        super().__init__()
         self.path = path
         self._client = client
         self._dirp = dirp
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def __del__(self):
-        try:
-            self.close()
-        except Exception:  # noqa: E722
-            # Best effort.
-            pass
-
     def __iter__(self) -> Iterator[DirEntry]:
         raise NotImplementedError()
 
-    def is_closed(self):
-        return not self._dirp
 
-    def close(self):
-        if self.is_closed():
-            return
-        self.closedir()
-        self._dirp = 0
-
-    def closedir(self):
-        raise NotImplementedError()
-
-
-class File:
+class File(Allocated):
     CHUNK_SIZE = 1024
 
     def __init__(self, client, fd: int):
         """
         :param rpcclient.client.client.Client client:
         :param fd:
         """
+        super().__init__()
         self._client = client
         self.fd = fd
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def close(self):
+    def _deallocate(self):
         """ close(fd) at remote. read man for more details. """
         fd = self._client.symbols.close(self.fd).c_int32
         if fd < 0:

src/rpcclient/rpcclient/network.py

@@ -3,31 +3,27 @@
 from collections import namedtuple
 
 from rpcclient.exceptions import BadReturnValueError
+from rpcclient.allocated import Allocated
 from rpcclient.structs.consts import AF_UNIX, AF_INET, SOCK_STREAM
 from rpcclient.structs.generic import sockaddr_in, sockaddr_un, ifaddrs, sockaddr, hostent
 
 Interface = namedtuple('Interface', 'name address netmask broadcast')
 Hostentry = namedtuple('Hostentry', 'name aliases addresses')
 
 
-class Socket:
+class Socket(Allocated):
     CHUNK_SIZE = 1024
 
     def __init__(self, client, fd: int):
         """
         :param rpcclient.client.client.Client client:
         :param fd:
         """
+        super().__init__()
         self._client = client
         self.fd = fd
 
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def close(self):
+    def _deallocate(self):
         """ close(fd) at remote. read man for more details. """
         fd = self._client.symbols.close(self.fd).c_int32
         if fd < 0: