UnsafeIPAddress: Implement __(g|s)etstate__ and to ensure proper (un)pickling

Missing attributes in instance created by pickle.load cause AttributeError in
second part of ipa-server-install --external-ca.

https://fedorahosted.org/freeipa/ticket/6385

Reviewed-By: Martin Basti <mbasti@redhat.com>
Reviewed-By: Jan Cholasta <jcholast@redhat.com>

1 files changed, 22 insertions, 0 deletions

diff --git a/ipapython/ipautil.py b/ipapython/ipautil.py
index b0aa26285..f7d75376b 100644
--- a/ipapython/ipautil.py
+++ b/ipapython/ipautil.py
@@ -125,6 +125,17 @@ class UnsafeIPAddress(netaddr.IPAddress):
         super(UnsafeIPAddress, self).__init__(addr,
                                               flags=self.netaddr_ip_flags)
 
+    def __getstate__(self):
+        state = {
+            '_net': self._net,
+            'super_state': super(UnsafeIPAddress, self).__getstate__(),
+        }
+        return state
+
+    def __setstate__(self, state):
+        super(UnsafeIPAddress, self).__setstate__(state['super_state'])
+        self._net = state['_net']
+
 
 class CheckedIPAddress(UnsafeIPAddress):
     """IPv4 or IPv6 address with additional constraints.
@@ -203,6 +214,17 @@ class CheckedIPAddress(UnsafeIPAddress):
 
         self.prefixlen = self._net.prefixlen
 
+    def __getstate__(self):
+        state = {
+            'prefixlen': self.prefixlen,
+            'super_state': super(CheckedIPAddress, self).__getstate__(),
+        }
+        return state
+
+    def __setstate__(self, state):
+        super(CheckedIPAddress, self).__setstate__(state['super_state'])
+        self.prefixlen = state['prefixlen']
+
     def is_network_addr(self):
         return self == self._net.network