scapy_service payload gen: allow template_code without size property

Signed-off-by: Anton Kiselev <anton.kisel@gmail.com>

1 files changed, 11 insertions, 5 deletions

diff --git a/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py b/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
index f889bb93..88514aa8 100755
--- a/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
+++ b/scripts/automation/trex_control_plane/stl/services/scapy_server/scapy_service.py
@@ -309,6 +309,11 @@ def parse_template_code(template_code):
     template_code = re.sub("[\s]", '', template_code) # remove spaces
     return bytearray.fromhex(template_code)
 
+def verify_payload_size(size):
+    assert(size != None)
+    if (size > (1<<20)): # 1Mb ought to be enough for anybody
+        raise ValueError('size is too large')
+
 def generate_bytes(bytes_definition):
     # accepts a bytes definition object
     # {generate: random_bytes or random_ascii, seed: <seed_number>, size: <size_bytes>}
@@ -317,20 +322,21 @@ def generate_bytes(bytes_definition):
     gen_type = bytes_definition.get('generate')
     if gen_type == None:
         return b64_to_bytes(bytes_definition['base64'])
+    elif gen_type == 'template_code':
+        code = parse_template_code(bytes_definition["template_code"])
+        bytes_size = int(bytes_definition.get('size') or len(code))
+        verify_payload_size(bytes_size)
+        return generate_bytes_from_template(bytes_size, code)
     else:
         bytes_size = int(bytes_definition['size']) # required
         seed = int(bytes_definition.get('seed') or 12345) # optional
-        if (bytes_size > (1<<20)): # 1Mb ought to be enough for anybody
-            raise ValueError('size is too large')
+        verify_payload_size(bytes_size)
         if gen_type == 'random_bytes':
             return generate_random_bytes(bytes_size, seed, 0, 0xFF)
         elif gen_type == 'random_ascii':
             return generate_random_bytes(bytes_size, seed, 0x20, 0x7E)
         elif gen_type == 'template':
             return generate_bytes_from_template(bytes_size, b64_to_bytes(bytes_definition["template_base64"]))
-        elif gen_type == 'template_code':
-            return generate_bytes_from_template(bytes_size, parse_template_code(bytes_definition["template_code"]))
-
 
 class ScapyException(Exception): pass
 class Scapy_service(Scapy_service_api):