From c202aee9049684831d68bf4a99c88da9d2a45700 Mon Sep 17 00:00:00 2001
From: Ad Schellevis <ad@opnsense.org>
Date: Tue, 12 Dec 2017 20:24:22 +0100
Subject: [PATCH] ipsec, upgrade vici lib to 5.5.3.
 https://github.com/opnsense/core/issues/1981

---
 src/opnsense/scripts/ipsec/vici/protocol.py | 22 +++++++++++++-----
 src/opnsense/scripts/ipsec/vici/session.py  | 25 +++++++++++++++++++--
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/opnsense/scripts/ipsec/vici/protocol.py b/src/opnsense/scripts/ipsec/vici/protocol.py
index 855a7b2e2..370229463 100644
--- a/src/opnsense/scripts/ipsec/vici/protocol.py
+++ b/src/opnsense/scripts/ipsec/vici/protocol.py
@@ -20,15 +20,25 @@ class Transport(object):
         self.socket.sendall(struct.pack("!I", len(packet)) + packet)
 
     def receive(self):
-        raw_length = self.socket.recv(self.HEADER_LENGTH)
+        raw_length = self._recvall(self.HEADER_LENGTH)
         length, = struct.unpack("!I", raw_length)
-        payload = self.socket.recv(length)
+        payload = self._recvall(length)
         return payload
 
     def close(self):
         self.socket.shutdown(socket.SHUT_RDWR)
         self.socket.close()
 
+    def _recvall(self, count):
+        """Ensure to read count bytes from the socket"""
+        data = b""
+        while len(data) < count:
+            buf = self.socket.recv(count - len(data))
+            if not buf:
+                raise socket.error('Connection closed')
+            data += buf
+        return data
+
 
 class Packet(object):
     CMD_REQUEST = 0         # Named request message
@@ -52,7 +62,7 @@ class Packet(object):
 
     @classmethod
     def _named_request(cls, request_type, request, message=None):
-        request = request.encode()
+        request = request.encode("UTF-8")
         payload = struct.pack("!BB", request_type, len(request)) + request
         if message is not None:
             return payload + message
@@ -95,12 +105,12 @@ class Message(object):
     @classmethod
     def serialize(cls, message):
         def encode_named_type(marker, name):
-            name = name.encode()
+            name = name.encode("UTF-8")
             return struct.pack("!BB", marker, len(name)) + name
 
         def encode_blob(value):
             if not isinstance(value, bytes):
-                value = str(value).encode()
+                value = str(value).encode("UTF-8")
             return struct.pack("!H", len(value)) + value
 
         def serialize_list(lst):
@@ -137,7 +147,7 @@ class Message(object):
     def deserialize(cls, stream):
         def decode_named_type(stream):
             length, = struct.unpack("!B", stream.read(1))
-            return stream.read(length).decode()
+            return stream.read(length).decode("UTF-8")
 
         def decode_blob(stream):
             length, = struct.unpack("!H", stream.read(2))
diff --git a/src/opnsense/scripts/ipsec/vici/session.py b/src/opnsense/scripts/ipsec/vici/session.py
index 283e3d13d..1383fa778 100644
--- a/src/opnsense/scripts/ipsec/vici/session.py
+++ b/src/opnsense/scripts/ipsec/vici/session.py
@@ -53,6 +53,14 @@ class Session(object):
         """
         return self.handler.streamed_request("terminate", "control-log", sa)
 
+    def redirect(self, sa):
+        """Redirect an IKE_SA.
+
+        :param sa: the SA to redirect
+        :type sa: dict
+        """
+        self.handler.request("redirect", sa)
+
     def install(self, policy):
         """Install a trap, drop or bypass policy defined by a CHILD_SA config.
 
@@ -158,6 +166,17 @@ class Session(object):
         """
         self.handler.request("load-shared", secret)
 
+    def flush_certs(self, filter=None):
+        """Flush the volatile certificate cache.
+
+        Flush the certificate stored temporarily in the cache. The filter
+        allows to flush only a certain type of certificates, e.g. CRLs.
+
+        :param filter: flush only certificates of a given type (optional)
+        :type filter: dict
+        """
+        self.handler.request("flush-certs", filter)
+
     def clear_creds(self):
         """Clear credentials loaded over vici.
 
@@ -189,13 +208,15 @@ class Session(object):
         """
         self.handler.request("unload-pool", pool_name)
 
-    def get_pools(self):
+    def get_pools(self, options):
         """Retrieve loaded pools.
 
+        :param options: filter by name and/or retrieve leases (optional)
+        :type options: dict
         :return: loaded pools
         :rtype: dict
         """
-        return self.handler.request("get-pools")
+        return self.handler.request("get-pools", options)
 
     def listen(self, event_types):
         """Register and listen for the given events.