vici: Move event (un-)registration to a helper method in Python library

Also make sure events are unregistered in case of exceptions in
streamed_request().

3 files changed, 60 insertions, 49 deletions

diff --git a/src/libcharon/plugins/vici/python/LICENSE b/src/libcharon/plugins/vici/python/LICENSE
index 111523ca8..54f2158dc 100644
--- a/src/libcharon/plugins/vici/python/LICENSE
+++ b/src/libcharon/plugins/vici/python/LICENSE
@@ -1,4 +1,6 @@
 Copyright (c) 2015 Björn Schuberg
+Copyright (c) 2015 Martin Willi
+Copyright (c) 2015 Tobias Brunner
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/src/libcharon/plugins/vici/python/vici/exception.py b/src/libcharon/plugins/vici/python/vici/exception.py
index 36384e556..757ac51a9 100644
--- a/src/libcharon/plugins/vici/python/vici/exception.py
+++ b/src/libcharon/plugins/vici/python/vici/exception.py
@@ -8,3 +8,6 @@ class SessionException(Exception):
 
 class CommandException(Exception):
     """Command result exception."""
+
+class EventUnknownException(Exception):
+    """Event unknown exception."""
diff --git a/src/libcharon/plugins/vici/python/vici/session.py b/src/libcharon/plugins/vici/python/vici/session.py
index dee58699d..53ce10dbd 100644
--- a/src/libcharon/plugins/vici/python/vici/session.py
+++ b/src/libcharon/plugins/vici/python/vici/session.py
@@ -1,7 +1,7 @@
 import collections
 import socket
 
-from .exception import SessionException, CommandException
+from .exception import SessionException, CommandException, EventUnknownException
 from .protocol import Transport, Packet, Message
 
 
@@ -215,6 +215,32 @@ class SessionHandler(object):
         self.transport.send(packet)
         return Packet.parse(self.transport.receive())
 
+    def _register_unregister(self, event_type, register):
+        """Register or unregister for the given event.
+
+        :param event_type: event to register
+        :type event_type: str
+        :param register: whether to register or unregister
+        :type register: bool
+        """
+        if register:
+            packet = Packet.register_event(event_type)
+        else:
+            packet = Packet.unregister_event(event_type)
+        response = self._communicate(packet)
+        if response.response_type == Packet.EVENT_UNKNOWN:
+            raise EventUnknownException(
+                "Unknown event type '{event}'".format(event=event_type)
+            )
+        elif response.response_type != Packet.EVENT_CONFIRM:
+            raise SessionException(
+                "Unexpected response type {type}, "
+                "expected '{confirm}' (EVENT_CONFIRM)".format(
+                    type=response.response_type,
+                    confirm=Packet.EVENT_CONFIRM,
+                )
+            )
+
     def request(self, command, message=None):
         """Send request with an optional message.
 
@@ -265,57 +291,37 @@ class SessionHandler(object):
         if message is not None:
             message = Message.serialize(message)
 
-        # subscribe to event stream
-        packet = Packet.register_event(event_stream_type)
-        response = self._communicate(packet)
-
-        if response.response_type != Packet.EVENT_CONFIRM:
-            raise SessionException(
-                "Unexpected response type {type}, "
-                "expected '{confirm}' (EVENT_CONFIRM)".format(
-                    type=response.response_type,
-                    confirm=Packet.EVENT_CONFIRM,
-                )
-            )
-
-        # issue command, and read any event messages
-        packet = Packet.request(command, message)
-        self.transport.send(packet)
-        exited = False
-        while True:
-            response = Packet.parse(self.transport.receive())
-            if response.response_type == Packet.EVENT:
-                if not exited:
-                    try:
-                        yield Message.deserialize(response.payload)
-                    except GeneratorExit:
-                        exited = True
-                        pass
+        self._register_unregister(event_stream_type, True);
+
+        try:
+            packet = Packet.request(command, message)
+            self.transport.send(packet)
+            exited = False
+            while True:
+                response = Packet.parse(self.transport.receive())
+                if response.response_type == Packet.EVENT:
+                    if not exited:
+                        try:
+                            yield Message.deserialize(response.payload)
+                        except GeneratorExit:
+                            exited = True
+                            pass
+                else:
+                    break
+
+            if response.response_type == Packet.CMD_RESPONSE:
+                command_response = Message.deserialize(response.payload)
             else:
-                break
-
-        if response.response_type == Packet.CMD_RESPONSE:
-            command_response = Message.deserialize(response.payload)
-        else:
-            raise SessionException(
-                "Unexpected response type {type}, "
-                "expected '{response}' (CMD_RESPONSE)".format(
-                    type=response.response_type,
-                    response=Packet.CMD_RESPONSE
+                raise SessionException(
+                    "Unexpected response type {type}, "
+                    "expected '{response}' (CMD_RESPONSE)".format(
+                        type=response.response_type,
+                        response=Packet.CMD_RESPONSE
+                    )
                 )
-            )
 
-        # unsubscribe from event stream
-        packet = Packet.unregister_event(event_stream_type)
-        response = self._communicate(packet)
-        if response.response_type != Packet.EVENT_CONFIRM:
-            raise SessionException(
-                "Unexpected response type {type}, "
-                "expected '{confirm}' (EVENT_CONFIRM)".format(
-                    type=response.response_type,
-                    confirm=Packet.EVENT_CONFIRM,
-                )
-            )
+        finally:
+            self._register_unregister(event_stream_type, False);
 
         # evaluate command result, if any
         if "success" in command_response: