extras/packetforge/ProtocolHeader.py

   1 # Copyright (c) 2022 Intel and/or its affiliates.
   2 # Licensed under the Apache License, Version 2.0 (the "License");
   3 # you may not use this file except in compliance with the License.
   4 # You may obtain a copy of the License at:
   5 #
   6 #     http://www.apache.org/licenses/LICENSE-2.0
   7 #
   8 # Unless required by applicable law or agreed to in writing, software
   9 # distributed under the License is distributed on an "AS IS" BASIS,
  10 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11 # See the License for the specific language governing permissions and
  12 # limitations under the License.
  13
  14 from ProtocolHeaderAttribute import *
  15 from ProtocolHeaderField import *
  16 from InputFormat import *
  17 import ExpressionConverter
  18 import copy
  19
  20
  21 class ProtocolHeader:
  22     def __init__(self, node):
  23         self.fields = []
  24         self.attributes = []
  25         self.fieldDict = {}
  26         self.attributeDict = {}
  27         self.Buffer = []
  28         self.Mask = []
  29
  30         self.node = node
  31         for field in self.node.fields:
  32             phf = ProtocolHeaderField(field.Size, field.DefaultValue, None, field)
  33             self.fields.append(phf)
  34             if field.Name != "reserved":
  35                 self.fieldDict[field.Name] = phf
  36
  37         for attr in self.node.attributes:
  38             pha = ProtocolHeaderAttribute(attr.Size, attr.DefaultValue, attr)
  39             self.attributes.append(pha)
  40             self.attributeDict[attr.Name] = pha
  41
  42     def Name(self):
  43         return self.node.Name
  44
  45     def Fields(self):
  46         return self.fields
  47
  48     def Attributes(self):
  49         return self.attributes
  50
  51     def setField(self, name, expression, auto):
  52         if name == "reserved":
  53             return False
  54
  55         if name not in self.fieldDict:
  56             return False
  57
  58         field = self.fieldDict[name]
  59
  60         if field.UpdateValue(expression, auto):
  61             field.UpdateSize()
  62             return True
  63
  64         return False
  65
  66     def SetField(self, name, expression):
  67         return self.setField(name, expression, False)
  68
  69     def SetFieldAuto(self, name, expression):
  70         return self.setField(name, expression, True)
  71
  72     def SetAttribute(self, name, expression):
  73         if name not in self.attributeDict:
  74             return False
  75         attr = self.attributeDict[name]
  76
  77         return attr.UpdateValue(expression)
  78
  79     def SetMask(self, name, expression):
  80         if name not in self.fieldDict:
  81             return False
  82         field = self.fieldDict[name]
  83
  84         return field.UpdateMask(expression)
  85
  86     def resolveOptional(self, condition):
  87         if condition == None:
  88             return True
  89
  90         tokens = condition.split("|")
  91
  92         if len(tokens) > 1:
  93             result = False
  94
  95             for token in tokens:
  96                 result |= self.resolveOptional(token)
  97
  98             return result
  99
 100         tokens = condition.split("&")
 101
 102         if len(tokens) > 1:
 103             result = True
 104
 105             for token in tokens:
 106                 result &= self.resolveOptional(token)
 107
 108             return result
 109
 110         key = None
 111         value = None
 112
 113         if "!=" in tokens[0]:
 114             index = tokens[0].find("!=")
 115             key = tokens[0][:index].strip()
 116             value = tokens[0][index + 1 :].strip()
 117         elif "=" in tokens[0]:
 118             index = tokens[0].find("=")
 119             key = tokens[0][:index].strip()
 120             value = tokens[0][index + 1 :].strip()
 121         else:
 122             return False
 123
 124         if key not in self.fieldDict:
 125             return False
 126
 127         f = self.fieldDict[key]
 128         return ExpressionConverter.Equal(f.Value, value)
 129
 130     def resolveSize(self, exp):
 131         shift = 0
 132         key = exp
 133
 134         if "<<" in exp:
 135             offset = exp.find("<<")
 136             key = exp[0:offset].strip()
 137             shift = int(exp[offset + 2 :].strip())
 138
 139         if self.fieldDict.has_key(key):
 140             field = self.fieldDict[key]
 141             _, u16 = ExpressionConverter.ToNum(field.Value)
 142             if u16:
 143                 return u16 << shift
 144             else:
 145                 return 0
 146
 147         if self.attributeDict.has_key(key):
 148             attr = self.attributeDict[key]
 149             _, u16 = ExpressionConverter.ToNum(attr.Value)
 150             if u16:
 151                 return u16 << shift
 152             else:
 153                 return 0
 154
 155         return 0
 156
 157     def Adjust(self):
 158         autoIncreases = []
 159         increaseHeaders = []
 160
 161         self.resolveAllSize()
 162
 163         for phf in self.fields:
 164             if phf.Field.IsAutoIncrease:
 165                 autoIncreases.append(phf)
 166             if phf.Field.IsIncreaseLength and self.resolveOptional(phf.Field.Optional):
 167                 increaseHeaders.append(phf)
 168
 169         for f1 in autoIncreases:
 170             for f2 in increaseHeaders:
 171                 f1.UpdateValue(
 172                     ExpressionConverter.IncreaseValue(f1.Value, f2.Size >> 3), True
 173                 )
 174
 175     def resolveAllSize(self):
 176         for phf in self.fields:
 177             if phf.Field.Optional != None and not self.resolveOptional(
 178                 phf.Field.Optional
 179             ):
 180                 size = 0
 181             else:
 182                 if phf.Field.VariableSize != None:
 183                     size = self.resolveSize(phf.Field.VariableSize)
 184                 else:
 185                     size = phf.Field.Size
 186             phf.Size = size
 187
 188     def GetSize(self):
 189         size = 0
 190
 191         for field in self.fields:
 192             size += field.Size
 193
 194         return size >> 3
 195
 196     def AppendAuto(self, size):
 197         for phf in self.fields:
 198             if not phf.Field.IsAutoIncrease:
 199                 continue
 200
 201             phf.UpdateValue(ExpressionConverter.IncreaseValue(phf.Value, size), True)
 202
 203     def getField(self, name):
 204         if not self.fieldDict.has_key(name):
 205             return None
 206         field = self.fieldDict[name]
 207
 208         return field.Value
 209
 210     def getAttribute(self, name):
 211         if not self.attributeDict.has_key(name):
 212             return None
 213
 214         return self.attributeDict[name].Value
 215
 216     def GetValue(self, name):
 217         result = self.getField(name)
 218
 219         if result == None:
 220             return self.getAttribute(name)
 221
 222         return result
 223
 224     def appendNum(self, big, exp, size):
 225         num = 0
 226         if exp != None:
 227             _, num = ExpressionConverter.ToNum(exp)
 228             if num == None:
 229                 print("Invalid byte expression")
 230                 return None
 231
 232         # cut msb
 233         num = num & ((1 << size) - 1)
 234         big = big << size
 235         big = big | num
 236         return big
 237
 238     def appendUInt64(self, big, exp, size):
 239         u64 = 0
 240         if exp != None:
 241             _, u64 = ExpressionConverter.ToNum(exp)
 242             if not u64:
 243                 print("Invalid UInt32 expression")
 244                 return False
 245
 246         # cut msb
 247         if size < 64:
 248             u64 = u64 & ((1 << size) - 1)
 249         big = big << size
 250         big = big | u64
 251         return big
 252
 253     def appendIPv4(self, big, exp):
 254         ipv4 = bytes(4)
 255         if exp != None:
 256             _, ipv4 = ExpressionConverter.ToIPv4Address(exp)
 257             if not ipv4:
 258                 print("Inavalid IPv4 Address")
 259                 return False
 260
 261         for i in range(len(ipv4)):
 262             big = big << 8
 263             big = big | ipv4[i]
 264
 265         return big
 266
 267     def appendIPv6(self, big, exp):
 268         ipv6 = bytes(16)
 269         if exp != None:
 270             _, ipv6 = ExpressionConverter.ToIPv6Address(exp)
 271             if not ipv6:
 272                 print("Inavalid IPv6 Address")
 273                 return False
 274
 275         for i in range(16):
 276             big = big << 8
 277             big = big | ipv6[i]
 278
 279         return big
 280
 281     def appendMAC(self, big, exp):
 282         mac = bytes(6)
 283         if exp != None:
 284             _, mac = ExpressionConverter.ToMacAddress(exp)
 285             if not mac:
 286                 print("Inavalid MAC Address")
 287                 return False
 288
 289         for i in range(6):
 290             big = big << 8
 291             big = big | mac[i]
 292
 293         return big
 294
 295     def appendByteArray(self, big, exp, size):
 296         array = bytes(size >> 3)
 297         if exp != None:
 298             _, array = ExpressionConverter.ToByteArray(exp)
 299             if not array:
 300                 print("Invalid byte array")
 301                 return False
 302
 303         for i in range(size >> 3):
 304             big = big << 8
 305             if i < len(array):
 306                 big = big | array[i]
 307
 308         return big
 309
 310     def append(self, big, phf):
 311         bigVal = big["bigVal"]
 312         bigMsk = big["bigMsk"]
 313
 314         if phf.Field.IsReserved:
 315             bigVal <<= phf.Size
 316             bigMsk <<= phf.Size
 317             big.update(bigVal=bigVal, bigMsk=bigMsk)
 318             return big, phf.Size
 319
 320         size = phf.Size
 321
 322         if (
 323             phf.Field.Format == InputFormat.u8
 324             or phf.Field.Format == InputFormat.u16
 325             or phf.Field.Format == InputFormat.u32
 326         ):
 327             bigVal = self.appendNum(bigVal, phf.Value, size)
 328             bigMsk = self.appendNum(bigMsk, phf.Mask, size)
 329
 330         elif phf.Field.Format == InputFormat.u64:
 331             bigVal = self.appendUInt64(bigVal, phf.Value, size)
 332             bigMsk = self.appendUInt64(bigMsk, phf.Mask, size)
 333
 334         elif phf.Field.Format == InputFormat.ipv4:
 335             bigVal = self.appendIPv4(bigVal, phf.Value)
 336             bigMsk = self.appendIPv4(bigMsk, phf.Mask)
 337
 338         elif phf.Field.Format == InputFormat.ipv6:
 339             bigVal = self.appendIPv6(bigVal, phf.Value)
 340             bigMsk = self.appendIPv6(bigMsk, phf.Mask)
 341
 342         elif phf.Field.Format == InputFormat.mac:
 343             bigVal = self.appendMAC(bigVal, phf.Value)
 344             bigMsk = self.appendMAC(bigMsk, phf.Mask)
 345
 346         elif phf.Field.Format == InputFormat.bytearray:
 347             bigVal = self.appendByteArray(bigVal, phf.Value, size)
 348             bigMsk = self.appendByteArray(bigMsk, phf.Mask, size)
 349
 350         else:
 351             print("Invalid input format")
 352
 353         big.update(bigVal=bigVal, bigMsk=bigMsk)
 354         return big, size
 355
 356     def Resolve(self):
 357         big = {"bigVal": 0, "bigMsk": 0}
 358         offset = 0
 359
 360         for phf in self.fields:
 361             if phf.Size == 0:
 362                 continue
 363
 364             big, bits = self.append(big, phf)
 365
 366             offset += bits
 367
 368         byteList1 = []
 369         byteList2 = []
 370
 371         bigVal = big["bigVal"]
 372         bigMsk = big["bigMsk"]
 373
 374         while offset > 0:
 375             byteList1.append(bigVal & 0xFF)
 376             byteList2.append(bigMsk & 0xFF)
 377             bigVal = bigVal >> 8
 378             bigMsk = bigMsk >> 8
 379             offset -= 8
 380
 381         byteList1.reverse()
 382         byteList2.reverse()
 383         buffer = copy.deepcopy(byteList1)
 384         mask = copy.deepcopy(byteList2)
 385
 386         self.Buffer = buffer
 387         self.Mask = mask