From 36d2ee150cb0515d65f67cb8657d1c93ce687074 Mon Sep 17 00:00:00 2001
From: "C. Gao" <gaoc@knownsec.com>
Date: Wed, 6 Jan 2016 11:23:07 +0800
Subject: [PATCH] ipset: fix comment and test

---
 salt/modules/ipset.py | 23 +++++++++++++++--------
 salt/states/ipset.py  | 13 ++++++++++---
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/salt/modules/ipset.py b/salt/modules/ipset.py
index a2ff9e8686cd..05bd67c7f53e 100644
--- a/salt/modules/ipset.py
+++ b/salt/modules/ipset.py
@@ -430,11 +430,13 @@ def check(set=None, entry=None, family='ipv4'):
     if isinstance(entry, list):
         entries = entry
     else:
-        if entry.find('-') != -1 and entry.count('-') == 1:
-            start, end = entry.split('-')
+        _entry = entry.split()[0]
+        _entry_extra = entry.split()[1:]
+        if _entry.find('-') != -1 and _entry.count('-') == 1:
+            start, end = _entry.split('-')
 
             if settype == 'hash:ip':
-                entries = [str(ipaddress.ip_address(ip)) for ip in long_range(
+                entries = [' '.join([str(ipaddress.ip_address(ip)), ' '.join(_entry_extra)]) for ip in long_range(
                     ipaddress.ip_address(start),
                     ipaddress.ip_address(end) + 1
                 )]
@@ -444,17 +446,22 @@ def check(set=None, entry=None, family='ipv4'):
                                                              ipaddress.ip_address(end))
                 entries = []
                 for network in networks:
-                    entries.append(network.with_prefixlen)
+                    _network = [str(ip) for ip in ipaddress.ip_network(network)]
+                    if len(_network) == 1:
+                        __network = ' '.join([str(_network[0]), ' '.join(_entry_extra)])
+                    else:
+                        __network = ' '.join([str(network), ' '.join(_entry_extra)])
+                    entries.append(__network)
             else:
                 entries = [entry]
 
-        elif entry.find('/') != -1 and entry.count('/') == 1:
+        elif _entry.find('/') != -1 and _entry.count('/') == 1:
             if settype == 'hash:ip':
-                entries = [str(ip) for ip in ipaddress.ip_network(entry)]
+                entries = [' '.join([str(ip), ' '.join(_entry_extra)]) for ip in ipaddress.ip_network(_entry)]
             elif settype == 'hash:net':
-                _entries = [str(ip) for ip in ipaddress.ip_network(entry)]
+                _entries = [str(ip) for ip in ipaddress.ip_network(_entry)]
                 if len(_entries) == 1:
-                    entries = [_entries[0]]
+                    entries = [' '.join([_entries[0], ' '.join(_entry_extra)])]
                 else:
                     entries = [entry]
             else:
diff --git a/salt/states/ipset.py b/salt/states/ipset.py
index ea5d4a733e34..aec721f65e31 100644
--- a/salt/states/ipset.py
+++ b/salt/states/ipset.py
@@ -184,6 +184,7 @@ def present(name, entry=None, family='ipv4', **kwargs):
            'changes': {},
            'result': None,
            'comment': ''}
+    test_flag = False
 
     if not entry:
         ret['result'] = False
@@ -198,19 +199,25 @@ def present(name, entry=None, family='ipv4', **kwargs):
 
     for entry in entries:
         _entry = '{0}'.format(entry)
-        if 'comment' in kwargs:
-            _entry = '{0} comment "{1}"'.format(entry, kwargs['comment'])
+        if 'timeout' in kwargs:
+            if 'comment' in _entry:
+                _entry = '{0} timeout {1} {2} {3}'.format(entry.split()[0], kwargs['timeout'], entry.split()[1], entry.split()[2])
+            else:
+                _entry = '{0} timeout {1}'.format(entry.split()[0], kwargs['timeout'])
 
         if __salt__['ipset.check'](kwargs['set_name'],
                                    _entry,
                                    family) is True:
-            ret['result'] = True
+            if test_flag is False:
+                ret['result'] = True
             ret['comment'] += 'entry for {0} already in set ({1}) for {2}\n'.format(
                 entry,
                 kwargs['set_name'],
                 family)
         else:
             if __opts__['test']:
+                test_flag = True
+                ret['result'] = None
                 ret['comment'] += 'entry {0} needs to be added to set {1} for family {2}\n'.format(
                     entry,
                     kwargs['set_name'],