vici: Add vici python protocol handler
[strongswan.git] / src / libcharon / plugins / vici / python / vici / protocol.py
1 import io
2 import socket
3 import struct
4
5 from collections import namedtuple
6
7 from .exception import DeserializationException
8
9
10 class Transport(object):
11 HEADER_LENGTH = 4
12 MAX_SEGMENT = 512 * 1024
13
14 def __init__(self, address="/var/run/charon.vici"):
15 self.address = address
16 self.socket = socket.socket(socket.AF_UNIX)
17 self.socket.connect(address)
18
19 def send(self, packet):
20 self.socket.sendall(struct.pack("!I", len(packet)) + packet)
21
22 def receive(self):
23 raw_length = self.socket.recv(self.HEADER_LENGTH)
24 length, = struct.unpack("!I", raw_length)
25 payload = self.socket.recv(length)
26 return payload
27
28 def close(self):
29 self.socket.shutdown(socket.SHUT_RDWR)
30 self.socket.close()
31
32
33 class Packet(object):
34 CMD_REQUEST = 0 # Named request message
35 CMD_RESPONSE = 1 # Unnamed response message for a request
36 CMD_UNKNOWN = 2 # Unnamed response if requested command is unknown
37 EVENT_REGISTER = 3 # Named event registration request
38 EVENT_UNREGISTER = 4 # Named event de-registration request
39 EVENT_CONFIRM = 5 # Unnamed confirmation for event (de-)registration
40 EVENT_UNKNOWN = 6 # Unnamed response if event (de-)registration failed
41 EVENT = 7 # Named event message
42
43 ParsedPacket = namedtuple(
44 "ParsedPacket",
45 ["response_type", "payload"]
46 )
47
48 ParsedEventPacket = namedtuple(
49 "ParsedEventPacket",
50 ["response_type", "event_type", "payload"]
51 )
52
53 @classmethod
54 def _named_request(cls, request_type, request, message=None):
55 payload = struct.pack("!BB", request_type, len(request)) + request
56 if message is not None:
57 return payload + message
58 else:
59 return payload
60
61 @classmethod
62 def request(cls, command, message=None):
63 return cls._named_request(cls.CMD_REQUEST, command, message)
64
65 @classmethod
66 def register_event(cls, event_type):
67 return cls._named_request(cls.EVENT_REGISTER, event_type)
68
69 @classmethod
70 def unregister_event(cls, event_type):
71 return cls._named_request(cls.EVENT_UNREGISTER, event_type)
72
73 @classmethod
74 def parse(cls, packet):
75 stream = FiniteStream(packet)
76 response_type, = struct.unpack("!B", stream.read(1))
77
78 if response_type == cls.EVENT:
79 length, = struct.unpack("!B", stream.read(1))
80 event_type = stream.read(length)
81 return cls.ParsedEventPacket(response_type, event_type, stream)
82 else:
83 return cls.ParsedPacket(response_type, stream)
84
85
86 class Message(object):
87 SECTION_START = 1 # Begin a new section having a name
88 SECTION_END = 2 # End a previously started section
89 KEY_VALUE = 3 # Define a value for a named key in the section
90 LIST_START = 4 # Begin a named list for list items
91 LIST_ITEM = 5 # Define an unnamed item value in the current list
92 LIST_END = 6 # End a previously started list
93
94 @classmethod
95 def serialize(cls, message):
96 def encode_named_type(marker, name):
97 name = str(name)
98 return struct.pack("!BB", marker, len(name)) + name
99
100 def encode_blob(value):
101 value = str(value)
102 return struct.pack("!H", len(value)) + value
103
104 def serialize_list(lst):
105 segment = str()
106 for item in lst:
107 segment += struct.pack("!B", cls.LIST_ITEM) + encode_blob(item)
108 return segment
109
110 def serialize_dict(d):
111 segment = str()
112 for key, value in d.iteritems():
113 if isinstance(value, dict):
114 segment += (
115 encode_named_type(cls.SECTION_START, key)
116 + serialize_dict(value)
117 + struct.pack("!B", cls.SECTION_END)
118 )
119 elif isinstance(value, list):
120 segment += (
121 encode_named_type(cls.LIST_START, key)
122 + serialize_list(value)
123 + struct.pack("!B", cls.LIST_END)
124 )
125 else:
126 segment += (
127 encode_named_type(cls.KEY_VALUE, key)
128 + encode_blob(value)
129 )
130 return segment
131
132 return serialize_dict(message)
133
134 @classmethod
135 def deserialize(cls, stream):
136 def decode_named_type(stream):
137 length, = struct.unpack("!B", stream.read(1))
138 return stream.read(length)
139
140 def decode_blob(stream):
141 length, = struct.unpack("!H", stream.read(2))
142 return stream.read(length)
143
144 def decode_list_item(stream):
145 marker, = struct.unpack("!B", stream.read(1))
146 while marker == cls.LIST_ITEM:
147 yield decode_blob(stream)
148 marker, = struct.unpack("!B", stream.read(1))
149
150 if marker != cls.LIST_END:
151 raise DeserializationException(
152 "Expected end of list at {pos}".format(pos=stream.tell())
153 )
154
155 section = {}
156 section_stack = []
157 while stream.has_more():
158 element_type, = struct.unpack("!B", stream.read(1))
159 if element_type == cls.SECTION_START:
160 section_name = decode_named_type(stream)
161 new_section = {}
162 section[section_name] = new_section
163 section_stack.append(section)
164 section = new_section
165
166 elif element_type == cls.LIST_START:
167 list_name = decode_named_type(stream)
168 section[list_name] = [item for item in decode_list_item(stream)]
169
170 elif element_type == cls.KEY_VALUE:
171 key = decode_named_type(stream)
172 section[key] = decode_blob(stream)
173
174 elif element_type == cls.SECTION_END:
175 if len(section_stack):
176 section = section_stack.pop()
177 else:
178 raise DeserializationException(
179 "Unexpected end of section at {pos}".format(
180 pos=stream.tell()
181 )
182 )
183
184 if len(section_stack):
185 raise DeserializationException("Expected end of section")
186 return section
187
188
189 class FiniteStream(io.BytesIO):
190 def __len__(self):
191 return len(self.getvalue())
192
193 def has_more(self):
194 return self.tell() < len(self)