OLD | NEW |
| (Empty) |
1 # Copyright (c) 2011 The Chromium Authors. All rights reserved. | |
2 # Use of this source code is governed by a BSD-style license that can be | |
3 # found in the LICENSE file. | |
4 | |
5 """Provides utility functions for TCP/UDP echo servers and clients. | |
6 | |
7 This program has classes and functions to encode, decode, calculate checksum | |
8 and verify the "echo request" and "echo response" messages. "echo request" | |
9 message is an echo message sent from the client to the server. "echo response" | |
10 message is a response from the server to the "echo request" message from the | |
11 client. | |
12 | |
13 The format of "echo request" message is | |
14 <version><checksum><payload_size><payload>. <version> is the version number | |
15 of the "echo request" protocol. <checksum> is the checksum of the <payload>. | |
16 <payload_size> is the size of the <payload>. <payload> is the echo message. | |
17 | |
18 The format of "echo response" message is | |
19 <version><checksum><payload_size><key><encoded_payload>.<version>, | |
20 <checksum> and <payload_size> are same as what is in the "echo request" message. | |
21 <encoded_payload> is encoded version of the <payload>. <key> is a randomly | |
22 generated key that is used to encode/decode the <payload>. | |
23 """ | |
24 | |
25 __author__ = 'rtenneti@google.com (Raman Tenneti)' | |
26 | |
27 | |
28 from itertools import cycle | |
29 from itertools import izip | |
30 import random | |
31 | |
32 | |
33 class EchoHeader(object): | |
34 """Class to keep header info of the EchoRequest and EchoResponse messages. | |
35 | |
36 This class knows how to parse the checksum, payload_size from the | |
37 "echo request" and "echo response" messages. It holds the checksum, | |
38 payload_size of the "echo request" and "echo response" messages. | |
39 """ | |
40 | |
41 # This specifies the version. | |
42 VERSION_STRING = '01' | |
43 | |
44 # This specifies the starting position of the checksum and length of the | |
45 # checksum. Maximum value for the checksum is less than (2 ** 31 - 1). | |
46 CHECKSUM_START = 2 | |
47 CHECKSUM_LENGTH = 10 | |
48 CHECKSUM_FORMAT = '%010d' | |
49 CHECKSUM_END = CHECKSUM_START + CHECKSUM_LENGTH | |
50 | |
51 # This specifies the starting position of the <payload_size> and length of the | |
52 # <payload_size>. Maximum number of bytes that can be sent in the <payload> is | |
53 # 9,999,999. | |
54 PAYLOAD_SIZE_START = CHECKSUM_END | |
55 PAYLOAD_SIZE_LENGTH = 7 | |
56 PAYLOAD_SIZE_FORMAT = '%07d' | |
57 PAYLOAD_SIZE_END = PAYLOAD_SIZE_START + PAYLOAD_SIZE_LENGTH | |
58 | |
59 def __init__(self, checksum=0, payload_size=0): | |
60 """Initializes the checksum and payload_size of self (EchoHeader). | |
61 | |
62 Args: | |
63 checksum: (int) | |
64 The checksum of the payload. | |
65 payload_size: (int) | |
66 The size of the payload. | |
67 """ | |
68 self.checksum = checksum | |
69 self.payload_size = payload_size | |
70 | |
71 def ParseAndInitialize(self, echo_message): | |
72 """Parses the echo_message and initializes self with the parsed data. | |
73 | |
74 This method extracts checksum, and payload_size from the echo_message | |
75 (echo_message could be either echo_request or echo_response messages) and | |
76 initializes self (EchoHeader) with checksum and payload_size. | |
77 | |
78 Args: | |
79 echo_message: (string) | |
80 The string representation of EchoRequest or EchoResponse objects. | |
81 Raises: | |
82 ValueError: Invalid data | |
83 """ | |
84 if not echo_message or len(echo_message) < EchoHeader.PAYLOAD_SIZE_END: | |
85 raise ValueError('Invalid data:%s' % echo_message) | |
86 self.checksum = int(echo_message[ | |
87 EchoHeader.CHECKSUM_START:EchoHeader.CHECKSUM_END]) | |
88 self.payload_size = int(echo_message[ | |
89 EchoHeader.PAYLOAD_SIZE_START:EchoHeader.PAYLOAD_SIZE_END]) | |
90 | |
91 def InitializeFromPayload(self, payload): | |
92 """Initializes the EchoHeader object with the payload. | |
93 | |
94 It calculates checksum for the payload and initializes self (EchoHeader) | |
95 with the calculated checksum and size of the payload. | |
96 | |
97 This method is used by the client code during testing. | |
98 | |
99 Args: | |
100 payload: (string) | |
101 The payload is the echo string (like 'hello'). | |
102 Raises: | |
103 ValueError: Invalid data | |
104 """ | |
105 if not payload: | |
106 raise ValueError('Invalid data:%s' % payload) | |
107 self.payload_size = len(payload) | |
108 self.checksum = Checksum(payload, self.payload_size) | |
109 | |
110 def __str__(self): | |
111 """String representation of the self (EchoHeader). | |
112 | |
113 Returns: | |
114 A string representation of self (EchoHeader). | |
115 """ | |
116 checksum_string = EchoHeader.CHECKSUM_FORMAT % self.checksum | |
117 payload_size_string = EchoHeader.PAYLOAD_SIZE_FORMAT % self.payload_size | |
118 return EchoHeader.VERSION_STRING + checksum_string + payload_size_string | |
119 | |
120 | |
121 class EchoRequest(EchoHeader): | |
122 """Class holds data specific to the "echo request" message. | |
123 | |
124 This class holds the payload extracted from the "echo request" message. | |
125 """ | |
126 | |
127 # This specifies the starting position of the <payload>. | |
128 PAYLOAD_START = EchoHeader.PAYLOAD_SIZE_END | |
129 | |
130 def __init__(self): | |
131 """Initializes EchoRequest object.""" | |
132 EchoHeader.__init__(self) | |
133 self.payload = '' | |
134 | |
135 def ParseAndInitialize(self, echo_request_data): | |
136 """Parses and Initializes the EchoRequest object from the echo_request_data. | |
137 | |
138 This method extracts the header information (checksum and payload_size) and | |
139 payload from echo_request_data. | |
140 | |
141 Args: | |
142 echo_request_data: (string) | |
143 The string representation of EchoRequest object. | |
144 Raises: | |
145 ValueError: Invalid data | |
146 """ | |
147 EchoHeader.ParseAndInitialize(self, echo_request_data) | |
148 if len(echo_request_data) <= EchoRequest.PAYLOAD_START: | |
149 raise ValueError('Invalid data:%s' % echo_request_data) | |
150 self.payload = echo_request_data[EchoRequest.PAYLOAD_START:] | |
151 | |
152 def InitializeFromPayload(self, payload): | |
153 """Initializes the EchoRequest object with payload. | |
154 | |
155 It calculates checksum for the payload and initializes self (EchoRequest) | |
156 object. | |
157 | |
158 Args: | |
159 payload: (string) | |
160 The payload string for which "echo request" needs to be constructed. | |
161 """ | |
162 EchoHeader.InitializeFromPayload(self, payload) | |
163 self.payload = payload | |
164 | |
165 def __str__(self): | |
166 """String representation of the self (EchoRequest). | |
167 | |
168 Returns: | |
169 A string representation of self (EchoRequest). | |
170 """ | |
171 return EchoHeader.__str__(self) + self.payload | |
172 | |
173 | |
174 class EchoResponse(EchoHeader): | |
175 """Class holds data specific to the "echo response" message. | |
176 | |
177 This class knows how to parse the "echo response" message. This class holds | |
178 key, encoded_payload and decoded_payload of the "echo response" message. | |
179 """ | |
180 | |
181 # This specifies the starting position of the |key_| and length of the |key_|. | |
182 # Minimum and maximum values for the |key_| are 100,000 and 999,999. | |
183 KEY_START = EchoHeader.PAYLOAD_SIZE_END | |
184 KEY_LENGTH = 6 | |
185 KEY_FORMAT = '%06d' | |
186 KEY_END = KEY_START + KEY_LENGTH | |
187 KEY_MIN_VALUE = 0 | |
188 KEY_MAX_VALUE = 999999 | |
189 | |
190 # This specifies the starting position of the <encoded_payload> and length | |
191 # of the <encoded_payload>. | |
192 ENCODED_PAYLOAD_START = KEY_END | |
193 | |
194 def __init__(self, key='', encoded_payload='', decoded_payload=''): | |
195 """Initializes the EchoResponse object.""" | |
196 EchoHeader.__init__(self) | |
197 self.key = key | |
198 self.encoded_payload = encoded_payload | |
199 self.decoded_payload = decoded_payload | |
200 | |
201 def ParseAndInitialize(self, echo_response_data=None): | |
202 """Parses and Initializes the EchoResponse object from echo_response_data. | |
203 | |
204 This method calls EchoHeader to extract header information from the | |
205 echo_response_data and it then extracts key and encoded_payload from the | |
206 echo_response_data. It holds the decoded payload of the encoded_payload. | |
207 | |
208 Args: | |
209 echo_response_data: (string) | |
210 The string representation of EchoResponse object. | |
211 Raises: | |
212 ValueError: Invalid echo_request_data | |
213 """ | |
214 EchoHeader.ParseAndInitialize(self, echo_response_data) | |
215 if len(echo_response_data) <= EchoResponse.ENCODED_PAYLOAD_START: | |
216 raise ValueError('Invalid echo_response_data:%s' % echo_response_data) | |
217 self.key = echo_response_data[EchoResponse.KEY_START:EchoResponse.KEY_END] | |
218 self.encoded_payload = echo_response_data[ | |
219 EchoResponse.ENCODED_PAYLOAD_START:] | |
220 self.decoded_payload = Crypt(self.encoded_payload, self.key) | |
221 | |
222 def InitializeFromEchoRequest(self, echo_request): | |
223 """Initializes EchoResponse with the data from the echo_request object. | |
224 | |
225 It gets the checksum, payload_size and payload from the echo_request object | |
226 and then encodes the payload with a random key. It also saves the payload | |
227 as decoded_payload. | |
228 | |
229 Args: | |
230 echo_request: (EchoRequest) | |
231 The EchoRequest object which has "echo request" message. | |
232 """ | |
233 self.checksum = echo_request.checksum | |
234 self.payload_size = echo_request.payload_size | |
235 self.key = (EchoResponse.KEY_FORMAT % | |
236 random.randrange(EchoResponse.KEY_MIN_VALUE, | |
237 EchoResponse.KEY_MAX_VALUE)) | |
238 self.encoded_payload = Crypt(echo_request.payload, self.key) | |
239 self.decoded_payload = echo_request.payload | |
240 | |
241 def __str__(self): | |
242 """String representation of the self (EchoResponse). | |
243 | |
244 Returns: | |
245 A string representation of self (EchoResponse). | |
246 """ | |
247 return EchoHeader.__str__(self) + self.key + self.encoded_payload | |
248 | |
249 | |
250 def Crypt(payload, key): | |
251 """Encodes/decodes the payload with the key and returns encoded payload. | |
252 | |
253 This method loops through the payload and XORs each byte with the key. | |
254 | |
255 Args: | |
256 payload: (string) | |
257 The string to be encoded/decoded. | |
258 key: (string) | |
259 The key used to encode/decode the payload. | |
260 | |
261 Returns: | |
262 An encoded/decoded string. | |
263 """ | |
264 return ''.join(chr(ord(x) ^ ord(y)) for (x, y) in izip(payload, cycle(key))) | |
265 | |
266 | |
267 def Checksum(payload, payload_size): | |
268 """Calculates the checksum of the payload. | |
269 | |
270 Args: | |
271 payload: (string) | |
272 The payload string for which checksum needs to be calculated. | |
273 payload_size: (int) | |
274 The number of bytes in the payload. | |
275 | |
276 Returns: | |
277 The checksum of the payload. | |
278 """ | |
279 checksum = 0 | |
280 length = min(payload_size, len(payload)) | |
281 for i in range (0, length): | |
282 checksum += ord(payload[i]) | |
283 return checksum | |
284 | |
285 | |
286 def GetEchoRequestData(payload): | |
287 """Constructs an "echo request" message from the payload. | |
288 | |
289 It builds an EchoRequest object from the payload and then returns a string | |
290 representation of the EchoRequest object. | |
291 | |
292 This is used by the TCP/UDP echo clients to build the "echo request" message. | |
293 | |
294 Args: | |
295 payload: (string) | |
296 The payload string for which "echo request" needs to be constructed. | |
297 | |
298 Returns: | |
299 A string representation of the EchoRequest object. | |
300 Raises: | |
301 ValueError: Invalid payload | |
302 """ | |
303 try: | |
304 echo_request = EchoRequest() | |
305 echo_request.InitializeFromPayload(payload) | |
306 return str(echo_request) | |
307 except (IndexError, ValueError): | |
308 raise ValueError('Invalid payload:%s' % payload) | |
309 | |
310 | |
311 def GetEchoResponseData(echo_request_data): | |
312 """Verifies the echo_request_data and returns "echo response" message. | |
313 | |
314 It builds the EchoRequest object from the echo_request_data and then verifies | |
315 the checksum of the EchoRequest is same as the calculated checksum of the | |
316 payload. If the checksums don't match then it returns None. It checksums | |
317 match, it builds the echo_response object from echo_request object and returns | |
318 string representation of the EchoResponse object. | |
319 | |
320 This is used by the TCP/UDP echo servers. | |
321 | |
322 Args: | |
323 echo_request_data: (string) | |
324 The string that echo servers send to the clients. | |
325 | |
326 Returns: | |
327 A string representation of the EchoResponse object. It returns None if the | |
328 echo_request_data is not valid. | |
329 Raises: | |
330 ValueError: Invalid echo_request_data | |
331 """ | |
332 try: | |
333 if not echo_request_data: | |
334 raise ValueError('Invalid payload:%s' % echo_request_data) | |
335 | |
336 echo_request = EchoRequest() | |
337 echo_request.ParseAndInitialize(echo_request_data) | |
338 | |
339 if Checksum(echo_request.payload, | |
340 echo_request.payload_size) != echo_request.checksum: | |
341 return None | |
342 | |
343 echo_response = EchoResponse() | |
344 echo_response.InitializeFromEchoRequest(echo_request) | |
345 | |
346 return str(echo_response) | |
347 except (IndexError, ValueError): | |
348 raise ValueError('Invalid payload:%s' % echo_request_data) | |
349 | |
350 | |
351 def DecodeAndVerify(echo_request_data, echo_response_data): | |
352 """Decodes and verifies the echo_response_data. | |
353 | |
354 It builds EchoRequest and EchoResponse objects from the echo_request_data and | |
355 echo_response_data. It returns True if the EchoResponse's payload and | |
356 checksum match EchoRequest's. | |
357 | |
358 This is used by the TCP/UDP echo clients for testing purposes. | |
359 | |
360 Args: | |
361 echo_request_data: (string) | |
362 The request clients sent to echo servers. | |
363 echo_response_data: (string) | |
364 The response clients received from the echo servers. | |
365 | |
366 Returns: | |
367 True if echo_request_data and echo_response_data match. | |
368 Raises: | |
369 ValueError: Invalid echo_request_data or Invalid echo_response | |
370 """ | |
371 | |
372 try: | |
373 echo_request = EchoRequest() | |
374 echo_request.ParseAndInitialize(echo_request_data) | |
375 except (IndexError, ValueError): | |
376 raise ValueError('Invalid echo_request:%s' % echo_request_data) | |
377 | |
378 try: | |
379 echo_response = EchoResponse() | |
380 echo_response.ParseAndInitialize(echo_response_data) | |
381 except (IndexError, ValueError): | |
382 raise ValueError('Invalid echo_response:%s' % echo_response_data) | |
383 | |
384 return (echo_request.checksum == echo_response.checksum and | |
385 echo_request.payload == echo_response.decoded_payload) | |
OLD | NEW |