From 7a2e95b3746faeb386c58dedbb40b71d81a5cff0 Mon Sep 17 00:00:00 2001
From: Garfield Freeman <gfreeman@paloaltonetworks.com>
Date: Tue, 9 Mar 2021 09:34:12 -0800
Subject: [PATCH] feat: Add rule hit count support

Fixes #239 - Adds `rule_hit_count()` to firewall.
---
 panos/base.py     |  2 +-
 panos/firewall.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++
 panos/policies.py |  3 ++
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/panos/base.py b/panos/base.py
index 09839ebb..62a5f68a 100644
--- a/panos/base.py
+++ b/panos/base.py
@@ -70,6 +70,7 @@ class PanObject(object):
     CHILDMETHODS = ()
     HA_SYNC = True
     TEMPLATE_NATIVE = False
+    _UNKNOWN_PANOS_VERSION = (sys.maxsize, 0, 0)
 
     def __init__(self, *args, **kwargs):
         # Set the 'name' variable
@@ -2257,7 +2258,6 @@ class VersionedPanObject(PanObject):
 
     """
 
-    _UNKNOWN_PANOS_VERSION = (sys.maxsize, 0, 0)
     _DEFAULT_NAME = None
     _TEMPLATE_DEVICE_XPATH = "/config/devices/entry[@name='localhost.localdomain']"
     _TEMPLATE_VSYS_XPATH = _TEMPLATE_DEVICE_XPATH + "/vsys/entry[@name='{vsys}']"
diff --git a/panos/firewall.py b/panos/firewall.py
index 1af2a199..1e9bc651 100644
--- a/panos/firewall.py
+++ b/panos/firewall.py
@@ -484,6 +484,90 @@ def organize_into_vsys(self, create_vsys_objects=True, refresh_vsys=True):
                             self.add(x)
                     break
 
+    def rule_hit_count(self, style=None, rules=None):
+        """Retrieve the rule hit count.
+
+        PAN-OS 8.1+
+
+        Args:
+            style (str): The rule style to use (used if the style cannot automatically
+                be determined from the `rules` param).  The style can be
+                "application-override", "authentication", "decryption", "dos",
+                "nat", "pbf", "qos", "sdwan", "security", or "tunnel-inspect".  If
+                unspecified and no rules are specified in which the type can be
+                inferred, then this param defaults to "security".
+            rules (list): A list of rules.  This can be a mix of `panos.policies`
+                instances or basic strings.  If no rules are given, then the hit
+                count for all rules is retrieved.
+
+        Returns:
+            dict:  A dict where the key is the rule name and the value is a dict of hit count information.
+
+        """
+        if self.retrieve_panos_version() < (8, 1, 0):
+            raise err.PanDeviceError("Rule hit count is supported in PAN-OS 8.1+")
+
+        vals = []
+        val_types = set([])
+        for r in rules or []:
+            if hasattr(r, "HIT_COUNT_STYLE") and hasattr(r, "uid"):
+                val_types.add(r.HIT_COUNT_STYLE)
+                vals.append(r.uid)
+            elif not isstring(r):
+                raise ValueError("Unsupported type sent: {0}".format(r.__class__))
+            else:
+                vals.append(r)
+
+        if len(val_types) > 1:
+            raise ValueError("Multiple types encountered: {0}".format(val_types))
+
+        if val_types:
+            style = val_types.pop()
+
+        if style is None:
+            style = "security"
+
+        cmd = ET.Element("show")
+        sub = ET.SubElement(cmd, "rule-hit-count")
+        sub = ET.SubElement(sub, "vsys")
+        sub = ET.SubElement(sub, "vsys-name")
+        sub = ET.SubElement(sub, "entry", {"name": self.vsys or "vsys1"})
+        sub = ET.SubElement(sub, "rule-base")
+        sub = ET.SubElement(sub, "entry", {"name": style})
+        sub = ET.SubElement(sub, "rules")
+
+        if not vals:
+            ET.SubElement(sub, "all")
+        else:
+            sub = ET.SubElement(sub, "list")
+            for r in vals:
+                ET.SubElement(sub, "member").text = r
+
+        res = self.op(ET.tostring(cmd, encoding="utf-8"), cmd_xml=False)
+
+        str_fields = ("latest",)
+        int_fields = (
+            "hit-count",
+            "last-hit-timestamp",
+            "last-reset-timestamp",
+            "first-hit-timestamp",
+            "rule-creation-timestamp",
+            "rule-modification-timestamp",
+        )
+        ans = {}
+        for elm in res.findall(
+            "./result/rule-hit-count/vsys/entry/rule-base/entry/rules/entry"
+        ):
+            name = elm.attrib["name"]
+            val = {"name": name}
+            for field in str_fields:
+                val[field] = elm.find("./{0}".format(field)).text
+            for field in int_fields:
+                val[field] = int(elm.find("./{0}".format(field)).text)
+            ans[name] = val
+
+        return ans
+
 
 class FirewallState(object):
     def __init__(self):
diff --git a/panos/policies.py b/panos/policies.py
index 18a0a454..ceff13f5 100644
--- a/panos/policies.py
+++ b/panos/policies.py
@@ -119,6 +119,7 @@ class SecurityRule(VersionedPanObject):
     # TODO: Add QoS variables
     SUFFIX = ENTRY
     ROOT = Root.VSYS
+    HIT_COUNT_STYLE = "security"
 
     def _setup(self):
         # xpaths
@@ -283,6 +284,7 @@ class NatRule(VersionedPanObject):
 
     SUFFIX = ENTRY
     ROOT = Root.VSYS
+    HIT_COUNT_STYLE = "nat"
 
     def _setup(self):
         # xpaths
@@ -613,6 +615,7 @@ class PolicyBasedForwarding(VersionedPanObject):
 
     SUFFIX = ENTRY
     ROOT = Root.VSYS
+    HIT_COUNT_STYLE = "pbf"
 
     def _setup(self):
         # xpaths