vici: Extract command wrappers in Python bindings
[strongswan.git] / src / libcharon / plugins / vici / python / vici / session.py
1 import collections
2 import socket
3
4 from .exception import SessionException, CommandException, EventUnknownException
5 from .protocol import Transport, Packet, Message
6 from .command_wrappers import CommandWrappers
7
8
9 class Session(CommandWrappers, object):
10 def __init__(self, sock=None):
11 if sock is None:
12 sock = socket.socket(socket.AF_UNIX)
13 sock.connect("/var/run/charon.vici")
14 self.transport = Transport(sock)
15
16 def _communicate(self, packet):
17 """Send packet over transport and parse response.
18
19 :param packet: packet to send
20 :type packet: :py:class:`vici.protocol.Packet`
21 :return: parsed packet in a tuple with message type and payload
22 :rtype: :py:class:`collections.namedtuple`
23 """
24 self.transport.send(packet)
25 return Packet.parse(self.transport.receive())
26
27 def _register_unregister(self, event_type, register):
28 """Register or unregister for the given event.
29
30 :param event_type: event to register
31 :type event_type: str
32 :param register: whether to register or unregister
33 :type register: bool
34 """
35 if register:
36 packet = Packet.register_event(event_type)
37 else:
38 packet = Packet.unregister_event(event_type)
39 response = self._communicate(packet)
40 if response.response_type == Packet.EVENT_UNKNOWN:
41 raise EventUnknownException(
42 "Unknown event type '{event}'".format(event=event_type)
43 )
44 elif response.response_type != Packet.EVENT_CONFIRM:
45 raise SessionException(
46 "Unexpected response type {type}, "
47 "expected '{confirm}' (EVENT_CONFIRM)".format(
48 type=response.response_type,
49 confirm=Packet.EVENT_CONFIRM,
50 )
51 )
52
53 def request(self, command, message=None):
54 """Send request with an optional message.
55
56 :param command: command to send
57 :type command: str
58 :param message: message (optional)
59 :type message: str
60 :return: command result
61 :rtype: dict
62 """
63 if message is not None:
64 message = Message.serialize(message)
65 packet = Packet.request(command, message)
66 response = self._communicate(packet)
67
68 if response.response_type != Packet.CMD_RESPONSE:
69 raise SessionException(
70 "Unexpected response type {type}, "
71 "expected '{response}' (CMD_RESPONSE)".format(
72 type=response.response_type,
73 response=Packet.CMD_RESPONSE
74 )
75 )
76
77 command_response = Message.deserialize(response.payload)
78 if "success" in command_response:
79 if command_response["success"] != b"yes":
80 raise CommandException(
81 "Command failed: {errmsg}".format(
82 errmsg=command_response["errmsg"]
83 )
84 )
85
86 return command_response
87
88 def streamed_request(self, command, event_stream_type, message=None):
89 """Send command request and collect and return all emitted events.
90
91 :param command: command to send
92 :type command: str
93 :param event_stream_type: event type emitted on command execution
94 :type event_stream_type: str
95 :param message: message (optional)
96 :type message: str
97 :return: generator for streamed event responses as dict
98 :rtype: generator
99 """
100 if message is not None:
101 message = Message.serialize(message)
102
103 self._register_unregister(event_stream_type, True)
104
105 try:
106 packet = Packet.request(command, message)
107 self.transport.send(packet)
108 exited = False
109 while True:
110 response = Packet.parse(self.transport.receive())
111 if response.response_type == Packet.EVENT:
112 if not exited:
113 try:
114 yield Message.deserialize(response.payload)
115 except GeneratorExit:
116 exited = True
117 pass
118 else:
119 break
120
121 if response.response_type == Packet.CMD_RESPONSE:
122 command_response = Message.deserialize(response.payload)
123 else:
124 raise SessionException(
125 "Unexpected response type {type}, "
126 "expected '{response}' (CMD_RESPONSE)".format(
127 type=response.response_type,
128 response=Packet.CMD_RESPONSE
129 )
130 )
131
132 finally:
133 self._register_unregister(event_stream_type, False)
134
135 # evaluate command result, if any
136 if "success" in command_response:
137 if command_response["success"] != b"yes":
138 raise CommandException(
139 "Command failed: {errmsg}".format(
140 errmsg=command_response["errmsg"]
141 )
142 )
143
144 def listen(self, event_types):
145 """Register and listen for the given events.
146
147 :param event_types: event types to register
148 :type event_types: list
149 :return: generator for streamed event responses as (event_type, dict)
150 :rtype: generator
151 """
152 for event_type in event_types:
153 self._register_unregister(event_type, True)
154
155 try:
156 while True:
157 response = Packet.parse(self.transport.receive())
158 if response.response_type == Packet.EVENT:
159 try:
160 msg = Message.deserialize(response.payload)
161 yield response.event_type, msg
162 except GeneratorExit:
163 break
164
165 finally:
166 for event_type in event_types:
167 self._register_unregister(event_type, False)