inet_diag: Allow passing a states mask to inet_diat.create()

And use it in pss to implement --listening/-l and --all/-a.

Signed-off-by: Arnaldo Carvalho de Melo <acme@redhat.com>

2 files changed, 23 insertions, 13 deletions

diff --git a/pss.py b/pss.py
index 9d33e4d..8e1103a 100755
--- a/pss.py
+++ b/pss.py
@@ -50,8 +50,8 @@ def print_ms_timer(s):
 
 	return rc
 
-def print_sockets(show_options):
-	idiag = inet_diag.create()
+def print_sockets(show_options, states):
+	idiag = inet_diag.create(states = states)
 	while True:
 		try:
 			s = idiag.get()
@@ -121,8 +121,10 @@ def main():
 		print str(err)
 		sys.exit(2)
 
+	states = inet_diag.default_states;
+
 	if not opts:
-		print_sockets(False)
+		print_sockets(False, states)
 		sys.exit(0)
 
 	for o, a in opts:
@@ -133,9 +135,9 @@ def main():
    		elif o in ( "-r", "--resolve"):
 			not_implemented(o)
    		elif o in ( "-a", "--all"):
-			not_implemented(o)
+			states = inet_diag.SS_ALL;
    		elif o in ( "-l", "--listening"):
-			not_implemented(o)
+			states = 1 << inet_diag.SS_LISTEN;
    		elif o in ( "-o", "--options"):
 			show_options = True
    		elif o in ( "-e", "--extended"):
@@ -172,7 +174,7 @@ def main():
 			usage()
 			return
 
-	print_sockets(show_options)
+	print_sockets(show_options, states)
 
 if __name__ == '__main__':
     main()
diff --git a/python-inet_diag/inet_diag.c b/python-inet_diag/inet_diag.c
index 5d80b66..e5b4bd5 100644
--- a/python-inet_diag/inet_diag.c
+++ b/python-inet_diag/inet_diag.c
@@ -75,10 +75,10 @@ static const char *sstate_name[] = {
 
 #define SS_ALL ((1 << SS_MAX) - 1)
 
-static int default_states = SS_ALL & ~((1 << SS_LISTEN) |
-				       (1 << SS_CLOSE) |
-				       (1 << SS_TIME_WAIT) |
-				       (1 << SS_SYN_RECV));
+static const int default_states = SS_ALL & ~((1 << SS_LISTEN) |
+					     (1 << SS_CLOSE) |
+					     (1 << SS_TIME_WAIT) |
+					     (1 << SS_SYN_RECV));
 
 static const char *tmr_name[] = {
 	"off",
@@ -379,13 +379,19 @@ static PyTypeObject inet_diag_type = {
 /* constructor */
 static char inet_diag_create__doc__[] =
 "create() -- creates a new inet_diag socket.";
-static PyObject *inet_diag__create(PyObject *args __unused)
+static PyObject *inet_diag__create(PyObject *mself __unused, PyObject *args,
+				   PyObject *keywds)
 {
+	int states = default_states;
+	static char *kwlist[] = { "states", };
 	struct inet_diag *self = PyObject_NEW(struct inet_diag,
 					      &inet_diag_type);
 	if (self == NULL)
 		return NULL;
 
+	if (!PyArg_ParseTupleAndKeywords(args, keywds, "|i", kwlist, &states))
+		goto out_err;
+
 	self->socket = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG);
 	if (self->socket < 0)
 		goto out_err;
@@ -405,7 +411,7 @@ static PyObject *inet_diag__create(PyObject *args __unused)
 		},
 		.r = {
 			.idiag_family = AF_INET,
-			.idiag_states = default_states,
+			.idiag_states = states,
 		},
 	};
 	struct iovec iov[1] = {
@@ -436,7 +442,7 @@ static struct PyMethodDef python_inet_diag__methods[] = {
 	{
 		.ml_name  = "create",
 		.ml_meth  = (PyCFunction)inet_diag__create,
-		.ml_flags = METH_VARARGS,
+		.ml_flags = METH_VARARGS | METH_KEYWORDS,
 		.ml_doc	  = inet_diag_create__doc__
 	},
 	{	.ml_name = NULL, },
@@ -458,4 +464,6 @@ PyMODINIT_FUNC initinet_diag(void)
 	PyModule_AddIntConstant(m, "SS_LAST_ACK",    SS_LAST_ACK);
 	PyModule_AddIntConstant(m, "SS_LISTEN",	     SS_LISTEN);
 	PyModule_AddIntConstant(m, "SS_CLOSING",     SS_CLOSING);
+	PyModule_AddIntConstant(m, "SS_ALL",	     SS_ALL);
+	PyModule_AddIntConstant(m, "default_states", default_states);
 }