From 00fdadc11b3189c70a8155e95dc29f5dc24d92c7 Mon Sep 17 00:00:00 2001
From: Ad Schellevis <ad@opnsense.org>
Date: Sat, 12 Aug 2017 13:53:20 +0200
Subject: [PATCH] extend configd call "configctl filter list states" with
 filter, limit and total number of states for
 https://github.com/opnsense/core/pull/1771

---
 src/opnsense/scripts/filter/list_states.py    | 20 +++++++++++++++++--
 .../conf/actions.d/actions_filter.conf        |  2 +-
 src/www/diag_states_summary.php               |  2 +-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/opnsense/scripts/filter/list_states.py b/src/opnsense/scripts/filter/list_states.py
index 305ca1224..9b57ee0c8 100755
--- a/src/opnsense/scripts/filter/list_states.py
+++ b/src/opnsense/scripts/filter/list_states.py
@@ -33,6 +33,7 @@ import subprocess
 import os
 import sys
 import ujson
+import argparse
 
 
 def parse_address(addr):
@@ -53,7 +54,14 @@ def parse_address(addr):
     return parse_result
 
 if __name__ == '__main__':
-    result = {'details': []}
+    # parse input arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output', help='output type [json/text]', default='json')
+    parser.add_argument('--filter', help='filter results', default='')
+    parser.add_argument('--limit', help='limit number of results', default='')
+    inputargs = parser.parse_args()
+
+    result = {'details': [], 'total_entries': 0}
     with tempfile.NamedTemporaryFile() as output_stream:
         subprocess.call(['/sbin/pfctl', '-s', 'state'], stdout=output_stream, stderr=open(os.devnull, 'wb'))
         output_stream.seek(0)
@@ -62,6 +70,14 @@ if __name__ == '__main__':
             for line in data.split('\n'):
                 parts = line.split()
                 if len(parts) >= 6:
+                    # count total number of state table entries
+                    result['total_entries'] += 1
+                    # apply filter when provided
+                    if inputargs.filter != "" and line.lower().find(inputargs.filter) == -1:
+                        continue
+                    # limit results
+                    if inputargs.limit.isdigit() and len(result['details']) >= int(inputargs.limit):
+                        continue
                     record = dict()
                     record['nat_addr'] = None
                     record['nat_port'] = None
@@ -91,7 +107,7 @@ if __name__ == '__main__':
     result['total'] = len(result['details'])
 
     # handle command line argument (type selection)
-    if len(sys.argv) > 1 and sys.argv[1] == 'json':
+    if inputargs.output == 'json':
         print(ujson.dumps(result))
     else:
         # output plain
diff --git a/src/opnsense/service/conf/actions.d/actions_filter.conf b/src/opnsense/service/conf/actions.d/actions_filter.conf
index 93e972375..bbb3cec7c 100644
--- a/src/opnsense/service/conf/actions.d/actions_filter.conf
+++ b/src/opnsense/service/conf/actions.d/actions_filter.conf
@@ -37,7 +37,7 @@ message:request content of pf %s table
 
 [list.states]
 command:/usr/local/opnsense/scripts/filter/list_states.py
-parameters: %s
+parameters: --filter=%s --limit=%s
 type:script_output
 message:request pf states
 
diff --git a/src/www/diag_states_summary.php b/src/www/diag_states_summary.php
index 1846a914b..8b9d84a76 100644
--- a/src/www/diag_states_summary.php
+++ b/src/www/diag_states_summary.php
@@ -78,7 +78,7 @@ $dstipinfo = array();
 $allipinfo = array();
 $pairipinfo = array();
 
-$states = json_decode(configd_run("filter list states json"), true);
+$states = json_decode(configd_run("filter list states"), true);
 if(isset($states['details'])) {
   foreach($states['details'] as $state) {
     if (isset($state['nat_addr']) && $states['direction'] == 'out') {