]>
Commit | Line | Data |
---|---|---|
8c089cdd | 1 | import collections |
61fb10c8 | 2 | import socket |
8c089cdd | 3 | |
cf4cefee | 4 | from .exception import SessionException, CommandException, EventUnknownException |
6a31a0f6 | 5 | from .protocol import Transport, Packet, Message |
e0f7da86 | 6 | from .command_wrappers import CommandWrappers |
6a31a0f6 BS |
7 | |
8 | ||
e0f7da86 | 9 | class Session(CommandWrappers, object): |
61fb10c8 MW |
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") | |
e0f7da86 | 14 | self.transport = Transport(sock) |
8c089cdd BS |
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) | |
a47e431b | 25 | return Packet.parse(self.transport.receive()) |
8c089cdd | 26 | |
cf4cefee TB |
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 | ||
8c089cdd | 53 | def request(self, command, message=None): |
90e16837 | 54 | """Send request with an optional message. |
8c089cdd BS |
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 | ||
90e16837 MW |
77 | command_response = Message.deserialize(response.payload) |
78 | if "success" in command_response: | |
2e74aa0a | 79 | if command_response["success"] != b"yes": |
90e16837 MW |
80 | raise CommandException( |
81 | "Command failed: {errmsg}".format( | |
82 | errmsg=command_response["errmsg"] | |
83 | ) | |
84 | ) | |
85 | ||
86 | return command_response | |
8c089cdd BS |
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 | |
a47e431b MW |
97 | :return: generator for streamed event responses as dict |
98 | :rtype: generator | |
8c089cdd | 99 | """ |
8c089cdd BS |
100 | if message is not None: |
101 | message = Message.serialize(message) | |
102 | ||
e0f7da86 | 103 | self._register_unregister(event_stream_type, True) |
cf4cefee TB |
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) | |
a47e431b | 123 | else: |
cf4cefee TB |
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 | ) | |
8c089cdd | 130 | ) |
8c089cdd | 131 | |
cf4cefee | 132 | finally: |
e0f7da86 | 133 | self._register_unregister(event_stream_type, False) |
b5d17e55 MW |
134 | |
135 | # evaluate command result, if any | |
136 | if "success" in command_response: | |
2e74aa0a | 137 | if command_response["success"] != b"yes": |
b5d17e55 MW |
138 | raise CommandException( |
139 | "Command failed: {errmsg}".format( | |
140 | errmsg=command_response["errmsg"] | |
141 | ) | |
142 | ) | |
b3d91cc6 TB |
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: | |
e0f7da86 TB |
160 | msg = Message.deserialize(response.payload) |
161 | yield response.event_type, msg | |
b3d91cc6 TB |
162 | except GeneratorExit: |
163 | break | |
164 | ||
165 | finally: | |
166 | for event_type in event_types: | |
167 | self._register_unregister(event_type, False) |