From 448305f2ef42c06f4478a8764d9e9bd8912419de Mon Sep 17 00:00:00 2001
From: Chris PeBenito <cpebenito@tresys.com>
Date: Wed, 12 Nov 2014 13:39:36 -0500
Subject: [PATCH] Add NodeconQuery functions for matching on IP version.

---
 setools/nodeconquery.py | 28 ++++++++++++++++++++++++++++
 tests/nodeconquery.py   |  8 ++++++++
 2 files changed, 36 insertions(+)

diff --git a/setools/nodeconquery.py b/setools/nodeconquery.py
index 5343c5e..044166d 100644
--- a/setools/nodeconquery.py
+++ b/setools/nodeconquery.py
@@ -22,6 +22,7 @@ except ImportError:
     pass
 
 import re
+from socket import AF_INET, AF_INET6
 
 from . import compquery
 from . import contextquery
@@ -33,6 +34,7 @@ class NodeconQuery(contextquery.ContextQuery):
 
     def __init__(self, policy,
                  net="", net_overlap=False,
+                 version=0,
                  user="", user_regex=False,
                  role="", role_regex=False,
                  type_="", type_regex=False,
@@ -59,6 +61,7 @@ class NodeconQuery(contextquery.ContextQuery):
         self.policy = policy
 
         self.set_network(net, overlap=net_overlap)
+        self.set_ip_version(version)
         self.set_user(user, regex=user_regex)
         self.set_role(role, regex=role_regex)
         self.set_type(type_, regex=type_regex)
@@ -99,6 +102,9 @@ class NodeconQuery(contextquery.ContextQuery):
                     if not net == self.network:
                         continue
 
+            if self.version and self.version != n.ip_version:
+                continue
+
             if not self._match_context(
                     n.context,
                     self.user,
@@ -147,3 +153,25 @@ class NodeconQuery(contextquery.ContextQuery):
                 self.network_overlap = opts[k]
             else:
                 raise NameError("Invalid name option: {0}".format(k))
+
+    def set_ip_version(self, version):
+        """
+        Set the criteria for matching the IP version.
+
+        Parameter:
+        version     The address family to match.  (socket.AF_INET for
+                    IPv4 or socket.AF_INET6 for IPv6)
+
+        Exceptions:
+        ValueError  Invalid address family number.
+        """
+
+        if version:
+            if not (version == AF_INET or version == AF_INET6):
+                raise ValueError(
+                    "The address family must be {0} for IPv4 or {1} for IPv6.".format(AF_INET, AF_INET6))
+
+            self.version = version
+
+        else:
+            self.version = None
diff --git a/tests/nodeconquery.py b/tests/nodeconquery.py
index 3d048a3..838b25f 100644
--- a/tests/nodeconquery.py
+++ b/tests/nodeconquery.py
@@ -17,6 +17,7 @@
 #
 import sys
 import unittest
+from socket import AF_INET6
 
 from setools import SELinuxPolicy
 from setools.nodeconquery import NodeconQuery
@@ -39,6 +40,13 @@ class NodeconQueryTest(unittest.TestCase):
 
         self.assertEqual(numrules, q_numrules)
 
+    def test_001_ip_version(self):
+        """Nodecon query with IP version match."""
+        q = NodeconQuery(self.p, version=AF_INET6)
+
+        nodecons = sorted(n.address for n in q.results())
+        self.assertListEqual(["1100::", "1110::"], nodecons)
+
     def test_020_user_exact(self):
         """Nodecon query with context user exact match"""
         q = NodeconQuery(self.p, user="user20", user_regex=False)