Ved-vampir | 0c7e2d4 | 2015-03-18 17:18:47 +0300 | [diff] [blame^] | 1 | #!/usr/bin/env python |
| 2 | """ Protocol class """ |
| 3 | |
| 4 | import re |
| 5 | import zlib |
| 6 | import json |
| 7 | import binascii |
| 8 | import logging |
| 9 | |
| 10 | import msgpack |
| 11 | |
| 12 | from logger import define_logger |
| 13 | |
| 14 | # protocol contains 2 type of packet: |
| 15 | # 1 - header, which contains template schema of counters |
| 16 | # 2 - body, which contains only values in order as in template |
| 17 | # it uses msgpack for optimization |
| 18 | # |
| 19 | # packet has format: |
| 20 | # begin_data_prefixSIZE\n\nDATAend_data_postfix |
| 21 | # packet part has format: |
| 22 | # SIZE\n\rDATA |
| 23 | # |
| 24 | # DATA use archivation |
| 25 | |
| 26 | |
| 27 | class PacketException(Exception): |
| 28 | """ Exceptions from Packet""" |
| 29 | pass |
| 30 | |
| 31 | |
| 32 | class Packet(object): |
| 33 | """ Class proceed packet by protocol""" |
| 34 | |
| 35 | prefix = "begin_data_prefix" |
| 36 | postfix = "end_data_postfix" |
| 37 | header_prefix = "template" |
| 38 | # other fields |
| 39 | # is_begin |
| 40 | # is_end |
| 41 | # crc |
| 42 | # data |
| 43 | # data_len |
| 44 | |
| 45 | def __init__(self): |
| 46 | # preinit |
| 47 | self.is_begin = False |
| 48 | self.is_end = False |
| 49 | self.crc = None |
| 50 | self.data = "" |
| 51 | self.data_len = None |
| 52 | self.srv_template = None |
| 53 | self.clt_template = None |
| 54 | self.tmpl_size = 0 |
| 55 | |
| 56 | |
| 57 | def new_packet(self, part): |
| 58 | """ New packet adding """ |
| 59 | # proceed packet |
| 60 | try: |
| 61 | # get size |
| 62 | local_size_s, _, part = part.partition("\n\r") |
| 63 | local_size = int(local_size_s) |
| 64 | |
| 65 | # find prefix |
| 66 | begin = part.find(self.prefix) |
| 67 | if begin != -1: |
| 68 | # divide data if something before begin and prefix |
| 69 | from_i = begin + len(self.prefix) |
| 70 | part = part[from_i:] |
| 71 | # reset flags |
| 72 | self.is_begin = True |
| 73 | self.is_end = False |
| 74 | self.data = "" |
| 75 | # get size |
| 76 | data_len_s, _, part = part.partition("\n\r") |
| 77 | self.data_len = int(data_len_s) |
| 78 | # get crc |
| 79 | crc_s, _, part = part.partition("\n\r") |
| 80 | self.crc = int(crc_s) |
| 81 | |
| 82 | # bad size? |
| 83 | if local_size != self.data_len: |
| 84 | raise PacketException("Part size error") |
| 85 | |
| 86 | # find postfix |
| 87 | end = part.find(self.postfix) |
| 88 | if end != -1: |
| 89 | # divide postfix |
| 90 | part = part[:end] |
| 91 | self.is_end = True |
| 92 | |
| 93 | self.data += part |
| 94 | # check if it is end |
| 95 | if self.is_end: |
| 96 | self.data = zlib.decompress(self.data) |
| 97 | if self.data_len != len(self.data): |
| 98 | raise PacketException("Total size error") |
| 99 | if binascii.crc32(self.data) != self.crc: |
| 100 | raise PacketException("CRC error") |
| 101 | |
| 102 | # check, if it is template |
| 103 | if self.data.startswith(self.header_prefix): |
| 104 | self.srv_template = self.data |
| 105 | # template is for internal use |
| 106 | return None |
| 107 | |
| 108 | # decode values list |
| 109 | vals = msgpack.unpackb(self.data) |
| 110 | dump = self.srv_template % tuple(vals) |
| 111 | return dump |
| 112 | else: |
| 113 | return None |
| 114 | |
| 115 | |
| 116 | except PacketException as e: |
| 117 | # if something wrong - skip packet |
| 118 | logger = logging.getLogger(__name__) |
| 119 | logger.warning("Packet skipped: %s", e) |
| 120 | self.is_begin = False |
| 121 | self.is_end = False |
| 122 | return None |
| 123 | |
| 124 | except TypeError: |
| 125 | # if something wrong - skip packet |
| 126 | logger = logging.getLogger(__name__) |
| 127 | logger.warning("Packet skipped: doesn't match schema") |
| 128 | self.is_begin = False |
| 129 | self.is_end = False |
| 130 | return None |
| 131 | |
| 132 | except: |
| 133 | # if something at all wrong - skip packet |
| 134 | logger = logging.getLogger(__name__) |
| 135 | logger.warning("Packet skipped: something is wrong") |
| 136 | self.is_begin = False |
| 137 | self.is_end = False |
| 138 | return None |
| 139 | |
| 140 | |
| 141 | @staticmethod |
| 142 | def create_packet(data, part_size): |
| 143 | """ Create packet divided by parts with part_size from data |
| 144 | No compression here """ |
| 145 | # prepare data |
| 146 | data_len = "%i\n\r" % len(data) |
| 147 | header = "%s%s%s\n\r" % (Packet.prefix, data_len, binascii.crc32(data)) |
| 148 | compact_data = zlib.compress(data) |
| 149 | packet = "%s%s%s" % (header, compact_data, Packet.postfix) |
| 150 | |
| 151 | partheader_len = len(data_len) |
| 152 | |
| 153 | beg = 0 |
| 154 | end = part_size - partheader_len |
| 155 | |
| 156 | result = [] |
| 157 | while beg < len(packet): |
| 158 | block = packet[beg:beg+end] |
| 159 | result.append(data_len + block) |
| 160 | beg += end |
| 161 | |
| 162 | return result |
| 163 | |
| 164 | |
| 165 | def create_packet_v2(self, data, part_size): |
| 166 | """ Create packet divided by parts with part_size from data |
| 167 | Compressed """ |
| 168 | result = [] |
| 169 | # create and add to result template header |
| 170 | if self.srv_template is None: |
| 171 | perf_string = json.dumps(data) |
| 172 | self.create_answer_template(perf_string) |
| 173 | template = self.header_prefix + self.srv_template |
| 174 | header = Packet.create_packet(template, part_size) |
| 175 | result.extend(header) |
| 176 | |
| 177 | vals = self.get_matching_value_list(data) |
| 178 | body = msgpack.packb(vals) |
| 179 | parts = Packet.create_packet(body, part_size) |
| 180 | result.extend(parts) |
| 181 | return result |
| 182 | |
| 183 | |
| 184 | def get_matching_value_list(self, data): |
| 185 | """ Get values in order server expect""" |
| 186 | vals = range(0, self.tmpl_size) |
| 187 | |
| 188 | try: |
| 189 | for node, groups in self.clt_template.items(): |
| 190 | for group, counters in groups.items(): |
| 191 | for counter, index in counters.items(): |
| 192 | if not isinstance(index, dict): |
| 193 | vals[index] = data[node][group][counter] |
| 194 | else: |
| 195 | for k, i in index.items(): |
| 196 | vals[i] = data[node][group][counter][k] |
| 197 | |
| 198 | return vals |
| 199 | |
| 200 | except (IndexError, KeyError): |
| 201 | logger = logging.getLogger(__name__) |
| 202 | logger.error("Data don't match last schema") |
| 203 | raise PacketException("Data don't match last schema") |
| 204 | |
| 205 | |
| 206 | |
| 207 | def create_answer_template(self, perf_string): |
| 208 | """ Create template for server to insert counter values |
| 209 | Return tuple of server and clien templates + number of replaces""" |
| 210 | replacer = re.compile(": [0-9]+\.?[0-9]*") |
| 211 | # replace all values by %s |
| 212 | finditer = replacer.finditer(perf_string) |
| 213 | # server not need know positions |
| 214 | self.srv_template = "" |
| 215 | # client need positions |
| 216 | clt_template = "" |
| 217 | beg = 0 |
| 218 | k = 0 |
| 219 | # this could be done better? |
| 220 | for match in finditer: |
| 221 | # define input place in server template |
| 222 | self.srv_template += perf_string[beg:match.start()] |
| 223 | self.srv_template += ": %s" |
| 224 | # define match number in client template |
| 225 | clt_template += perf_string[beg:match.start()] |
| 226 | clt_template += ": %i" % k |
| 227 | |
| 228 | beg = match.end() |
| 229 | k += 1 |
| 230 | |
| 231 | # add tail |
| 232 | self.srv_template += perf_string[beg:] |
| 233 | clt_template += perf_string[beg:] |
| 234 | |
| 235 | self.tmpl_size = k |
| 236 | self.clt_template = json.loads(clt_template) |
| 237 | |
| 238 | |
| 239 | |
| 240 | define_logger(__name__) |