From 15e55af88e6104608487c1138641fcff54594d89 Mon Sep 17 00:00:00 2001
From: Nicolas Fort <nicolasfort1988@gmail.com>
Date: Sat, 23 Apr 2022 12:05:57 +0000
Subject: Firewall: T990: Modifications for new connection-status cli

---
 .../include/firewall/common-rule.xml.i             | 39 +++++++++++++---------
 python/vyos/firewall.py                            |  7 +++-
 smoketest/scripts/cli/test_firewall.py             |  5 +--
 3 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/interface-definitions/include/firewall/common-rule.xml.i b/interface-definitions/include/firewall/common-rule.xml.i
index e74ce4ee4..85e586e1b 100644
--- a/interface-definitions/include/firewall/common-rule.xml.i
+++ b/interface-definitions/include/firewall/common-rule.xml.i
@@ -95,25 +95,32 @@
     </constraint>
   </properties>
 </leafNode>
-<leafNode name="connection-status">
+<node name="connection-status">
   <properties>
     <help>Connection status</help>
-    <completionHelp>
-      <list>dnat snat</list>
-    </completionHelp>
-    <valueHelp>
-      <format>dnat</format>
-      <description>Match connections that are subject to destination NAT</description>
-    </valueHelp>
-    <valueHelp>
-      <format>snat</format>
-      <description>Match connections that are subject to source NAT</description>
-    </valueHelp>
-    <constraint>
-      <regex>^(dnat|snat)$</regex>
-    </constraint>
   </properties>
-</leafNode>
+  <children>
+    <leafNode name="nat">
+      <properties>
+        <help>NAT connection status</help>
+        <completionHelp>
+          <list>destination source</list>
+        </completionHelp>
+        <valueHelp>
+          <format>destination</format>
+          <description>Match connections that are subject to destination NAT</description>
+        </valueHelp>
+        <valueHelp>
+          <format>source</format>
+          <description>Match connections that are subject to source NAT</description>
+        </valueHelp>
+        <constraint>
+          <regex>^(destination|source)$</regex>
+        </constraint>
+      </properties>
+    </leafNode>
+  </children>
+</node>
 <leafNode name="protocol">
   <properties>
     <help>Protocol to match (protocol name, number, or "all")</help>
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index 06731dd64..04fd44173 100644
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -51,7 +51,12 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
 
     if 'connection_status' in rule_conf and rule_conf['connection_status']:
         status = rule_conf['connection_status']
-        output.append(f'ct status {{{status}}}')
+        if status['nat'] == 'destination':
+            nat_status = '{dnat}'
+            output.append(f'ct status {nat_status}')
+        if status['nat'] == 'source':
+            nat_status = '{snat}'
+            output.append(f'ct status {nat_status}')
 
     if 'protocol' in rule_conf and rule_conf['protocol'] != 'all':
         proto = rule_conf['protocol']
diff --git a/smoketest/scripts/cli/test_firewall.py b/smoketest/scripts/cli/test_firewall.py
index 13bf02cd1..f72dfb1f4 100755
--- a/smoketest/scripts/cli/test_firewall.py
+++ b/smoketest/scripts/cli/test_firewall.py
@@ -171,11 +171,12 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '2', 'state', 'invalid', 'enable'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'action', 'accept'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'state', 'new', 'enable'])
-        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'connection-status', 'dnat'])
+
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'connection-status', 'nat', 'destination'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'action', 'accept'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'state', 'new', 'enable'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'state', 'established', 'enable'])
-        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'connection-status', 'snat'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'connection-status', 'nat', 'source'])
 
         self.cli_set(['interfaces', 'ethernet', 'eth0', 'firewall', 'in', 'name', 'smoketest'])
 
-- 
cgit v1.2.3