This commit is contained in:
Ad Schellevis 2017-12-12 20:24:22 +01:00
parent 651cbdaae3
commit c202aee904
2 changed files with 39 additions and 8 deletions

View File

@ -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))

View File

@ -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.