inet_diag: Add inode() method

To test it also introduce the psk utility:

$ psk pidgin firefox

15558: pidgin
  ESTAB      0      0     192.168.1.120:60373    209.85.163.125:5222
  ESTAB      0      0     192.168.1.120:35224    207.46.124.219:1863

15877: /usr/lib64/firefox-3.0.8/firefox
  ESTAB      0      0     192.168.1.120:48280       192.168.1.1:80
  ESTAB      0      0     192.168.1.120:48363       192.168.1.1:80
  ESTAB      0      0     192.168.1.120:48217       192.168.1.1:80
$

It gets a snapshot of /proc/pids using procfs.pidstats, and creates a
dictionary of inet_diag.inet_sock objects keyed by inode, then traverses the
list of processes looking for /proc/PID/fd/socket:[] symlinks, correlating
processes to sockets.

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

4 files changed, 142 insertions, 0 deletions

diff --git a/MANIFEST b/MANIFEST
index fed2b55..d64c7c7 100644
--- a/MANIFEST
+++ b/MANIFEST
@@ -1,5 +1,6 @@
 COPYING
 pss.py
+psk.py
 python-inet_diag/inet_diag_copy.h
 python-inet_diag/inet_diag.c
 setup.py
diff --git a/psk.py b/psk.py
new file mode 100644
index 0000000..42d6574
--- /dev/null
+++ b/psk.py
@@ -0,0 +1,126 @@
+#! /usr/bin/python
+# -*- python -*-
+# -*- coding: utf-8 -*-
+#   Copyright (C) 2009 Red Hat Inc.
+#   Written by Arnaldo Carvalho de Melo <acme@redhat.com>
+#
+#   This application is free software; you can redistribute it and/or
+#   modify it under the terms of the GNU General Public License
+#   as published by the Free Software Foundation; version 2.
+#
+#   This application is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+#   General Public License for more details.
+
+import getopt, inet_diag, os, re, sys, procfs
+
+version="0.1"
+
+state_width = 10
+addr_width = 20
+
+def load_sockets():
+	idiag = inet_diag.create()
+	inodes = {}
+	while True:
+		try:
+			s = idiag.get()
+		except:
+			break
+		inodes[s.inode()] = s
+	return inodes
+
+ps = None
+inode_re = None
+inodes = None
+
+def thread_mapper(s):
+	global ps
+	try:
+		return [ int(s), ]
+	except:
+		pass
+	if not ps:
+		ps = procfs.pidstats()
+
+	try:
+		return ps.find_by_regex(re.compile(fnmatch.translate(s)))
+	except:
+		return ps.find_by_name(s)
+
+def print_sockets(pid, indent = 0):
+	header_printed = False
+	dirname = "/proc/%d/fd" % pid
+	for filename in os.listdir(dirname):
+		pathname = os.path.join(dirname, filename)
+		try:
+			linkto = os.readlink(pathname)
+		except:
+			continue
+		inode_match = inode_re.match(linkto)
+		if not inode_match:
+			continue
+		inode = int(inode_match.group(1))
+		if not inodes.has_key(inode):
+			continue
+		if not header_printed:
+			try:
+				print "\n%s%d: %s" % (indent * "  ", pid,
+						      procfs.process_cmdline(ps[pid]))
+			except:
+				return
+			header_printed = True
+		s = inodes[inode]
+		print "  %-*s %-6d %-6d %*s:%-5d %*s:%-5d" % \
+		      (state_width, s.state(),
+		       s.receive_queue(), s.write_queue(),
+		       addr_width, s.saddr(), s.sport(),
+		       addr_width, s.daddr(), s.dport())
+
+def usage():
+	print '''Usage: psk [ OPTIONS ]
+       psk [ OPTIONS ] [PID-LIST]
+   -h, --help		this message
+   -V, --version	output version information'''
+
+def main():
+	global ps, inode_re, inodes
+
+	try:
+		opts, args = getopt.getopt(sys.argv[1:],
+					   "hV",
+					   ("help", "version"))
+	except getopt.GetoptError, err:
+		usage()
+		print str(err)
+		sys.exit(2)
+
+	for o, a in opts:
+   		if o in ( "-V", "--version"):
+			print version
+			return
+		else:
+			usage()
+			return
+
+	if args:
+		pid_list = reduce(lambda i, j: i + j,
+                                  map(thread_mapper, args))
+				
+	inodes = load_sockets()
+	ps = procfs.pidstats()
+	pids = ps.keys()
+	pids.sort()
+	inode_re = re.compile(r"socket:\[(\d+)\]")
+	ps.reload_threads()
+	for pid in pids:
+		if args and pid not in pid_list:
+			continue
+		print_sockets(pid)
+		if ps[pid].has_key("threads"):
+			for tid in ps[pid]["threads"].keys():
+				print_sockets(tid, 1)
+
+if __name__ == '__main__':
+    main()
diff --git a/python-inet_diag/inet_diag.c b/python-inet_diag/inet_diag.c
index c8bd2d8..db41ac5 100644
--- a/python-inet_diag/inet_diag.c
+++ b/python-inet_diag/inet_diag.c
@@ -165,6 +165,14 @@ static PyObject *inet_socket__write_queue(struct inet_socket *self,
 	return Py_BuildValue("i", self->msg.idiag_wqueue);
 }
 
+static char inet_socket__inode_doc__[] =
+"inode() -- get internet socket associated inode";
+static PyObject *inet_socket__inode(struct inet_socket *self,
+				    PyObject *args __unused)
+{
+	return Py_BuildValue("i", self->msg.idiag_inode);
+}
+
 static char inet_socket__state_doc__[] =
 "sport() -- get internet socket state";
 static PyObject *inet_socket__state(struct inet_socket *self,
@@ -205,6 +213,11 @@ static struct PyMethodDef inet_socket__methods[] = {
 		.ml_doc	  = inet_socket__write_queue_doc__,
 	},
 	{
+		.ml_name  = "inode",
+		.ml_meth  = (PyCFunction)inet_socket__inode,
+		.ml_doc	  = inet_socket__inode_doc__,
+	},
+	{
 		.ml_name  = "state",
 		.ml_meth  = (PyCFunction)inet_socket__state,
 		.ml_doc	  = inet_socket__state_doc__,
diff --git a/rpm/SPECS/python-inet_diag.spec b/rpm/SPECS/python-inet_diag.spec
index 6b6b078..8bcf3bf 100644
--- a/rpm/SPECS/python-inet_diag.spec
+++ b/rpm/SPECS/python-inet_diag.spec
@@ -27,6 +27,7 @@ rm -rf %{buildroot}
 %{__python} setup.py install --skip-build --root %{buildroot}
 mkdir -p %{buildroot}%{_sbindir}
 cp -p pss.py %{buildroot}%{_sbindir}/pss
+cp -p psk.py %{buildroot}%{_sbindir}/psk
 
 %clean
 rm -rf %{buildroot}
@@ -35,6 +36,7 @@ rm -rf %{buildroot}
 %defattr(-,root,root)
 %doc COPYING
 %{_sbindir}/pss
+%{_sbindir}/psk
 %{python_sitearch}/inet_diag.so
 %if "%{python_ver}" >= "2.5"
 %{python_sitearch}/*.egg-info