Merge pull request #834 from Ostorlab/fix/fix_open_ports_parsing_from…

…_yaml_definition

Fix open ports parsing from yaml.

src/ostorlab/agent/definitions.py

@@ -54,7 +54,13 @@ def from_yaml(cls, file: io.TextIOWrapper) -> "AgentDefinition":
             mounts=definition.get("mounts", []),
             restart_policy=definition.get("restart_policy", ""),
             mem_limit=definition.get("mem_limit"),
-            open_ports=definition.get("open_ports", []),
+            open_ports=[
+                defintions.PortMapping(
+                    source_port=p.get("src_port"),
+                    destination_port=p.get("dest_port"),
+                )
+                for p in definition.get("open_ports", [])
+            ],
             restrictions=definition.get("restrictions", []),
             version=definition.get("version"),
             description=definition.get("description"),

tests/agent/definitions_test.py

@@ -3,6 +3,7 @@
 import io
 
 from ostorlab.agent import definitions
+from ostorlab.utils import defintions as utils_defintions
 
 
 def testAgentDefinitionFromYaml_whenYamlIsValid_returnsValidAgentDefinition():
@@ -42,3 +43,45 @@ def testAgentDefinitionFromYaml_whenYamlIsValid_returnsValidAgentDefinition():
     assert agent_definition.name == "agent1"
     assert agent_definition.in_selectors == ["in_selector1", "in_selector2"]
     assert agent_definition.out_selectors == ["out_selector1", "out_selector2"]
+
+
+def testAgentDefinitionFromYaml_withOpenPorts_correctlyParseTheSourceAndDestinationPorts() -> (
+    None
+):
+    """Ensure the source & destination values of the open ports fields are parsed & handled correctly."""
+    valid_yaml_data = """
+            kind: Agent
+            name: "agent42"
+            description: "Agent 42, not 41."
+            source: "https://github.com/"
+            in_selectors: 
+            - "in_selector1"
+            - "in_selector2"
+            out_selectors:
+            - "out_selector1"
+            - "out_selector2"
+            args:
+              - name: "template_urls"
+                type: "array"
+                description: "list of template urls to run."
+                value:
+                  - 'https://google.com'
+                  - 1
+            caps: [NET_ADMIN]
+            open_ports:
+              - src_port: 4242
+                dest_port: 4242
+        """
+
+    yaml_data_file = io.StringIO(valid_yaml_data)
+
+    agent_definition = definitions.AgentDefinition.from_yaml(yaml_data_file)
+
+    assert agent_definition.name == "agent42"
+    assert agent_definition.caps == ["NET_ADMIN"]
+    assert agent_definition.open_ports == [
+        utils_defintions.PortMapping(
+            source_port=4242,
+            destination_port=4242,
+        )
+    ]