import pcap
import sys
import os
import string
import time
import socket
import struct

ETHERTYPE_IP = '\x08\x00'

def decode_mac_packet(s):
	d = {}
	d['dst_mac'] = s[0:6]
	d['src_mac'] = s[6:12]
	d['ethertype'] = s[12:14]
	d['payload'] = s[14:]
	return d

def decode_ip_packet(s):
	d = {}
	tmp1, d['tos'], total_len, id, tmp2, d['ttl'], d['protocol'], checksum, src_ip, dst_ip = struct.unpack('BBHHHBBHII', s[:20])
	tmp2 = socket.ntohs(tmp2)
	d['version'] = (tmp1 & 0xf0) >> 4
	d['header_len'] = tmp1 & 0x0f
	d['total_len'] = socket.ntohs(total_len)
	d['id'] = socket.ntohs(id)
	d['flags'] = (tmp2 & 0xe000) >> 5
	d['fragment_offset'] = (tmp2 & 0x1fff)
	d['checksum'] = socket.ntohs(checksum)
	d['src_ip'] = src_ip
	d['dst_ip'] = dst_ip
	if d['header_len'] > 5:
		d['options'] = s[20:4*(d['header_len']-5)]
	else:
		d['options'] = None
	d['payload'] = s[4*d['header_len']:]
	return d

def decode_tcp_packet(s):
	d = {}
	src_port, dst_port, d['seq_num'], d['ack_num'], tmp1, tmp2, d['win_size'], d['checksum'], d['urgent_ptr'] = struct.unpack('HHIIBBHHH', s[:20])
	d['src_port'] = socket.ntohs(src_port)
	d['dst_port'] = socket.ntohs(dst_port)
	d['data_offset'], d['reserved'] = (((tmp1 & 0xf0) >> 4), (tmp1 & 0x0f))
	d['payload'] = s[d['data_offset'] << 2:]
	return d

last_print = time.time()
print_interval = 1.0
counters = {}
def print_packet(pktlen, data, timestamp):
	global last_print, counters
	if not data:
		return

	#print 'MAC',
	mac_data = decode_mac_packet(data)
	if mac_data['ethertype'] != ETHERTYPE_IP:
		#print '!IP'
		return

	#print 'IP',
	ip_data = decode_ip_packet(mac_data['payload'])
	if ip_data['protocol'] != socket.IPPROTO_TCP:
		#print '!TCP'
		return

	#print 'TCP',
	tcp_data = decode_tcp_packet(ip_data['payload'])

	inode = lookup_tcp_socket_inode(ip_data['src_ip'], tcp_data['src_port'], ip_data['dst_ip'], tcp_data['dst_port'])
	pid, program = lookup_socket_inode_owner(inode)
	
	counters.setdefault(pid, {'name': program, 'count': 0})['count'] += pktlen
	#print inode, pid, program, pktlen
	#print pid, program, pktlen
	if time.time() - last_print > print_interval:
		os.system('clear')
		for pid, info in counters.iteritems():
			print '%5s %20s: %.2fkb/s' % (pid, info['name'], info['count']/1024.0/print_interval)
		last_print = time.time()
		counters = {}

def lookup_tcp_socket_inode(src_ip, src_port, dst_ip, dst_port):
	src_addr = '%08X:%04X' % (src_ip, src_port)
	dst_addr = '%08X:%04X' % (dst_ip, dst_port)
	inode = None

	#print src_addr, dst_addr,

	f = file('/proc/net/tcp', 'rb')
	for line in f:
		fields = line.split()
		if (fields[1] == src_addr and fields[2] == dst_addr) or (fields[1] == dst_addr and fields[2] == src_addr):
			inode = fields[9]
			break
	f.close()

	return inode

def lookup_socket_inode_owner(inode):
	for pid in os.listdir('/proc/'):
		try:
			int(pid, 10)
		except:
			continue
		fd_dir = '/proc/' + pid + '/fd/'
		try:
			for fd in os.listdir(fd_dir):
				path = os.readlink(fd_dir + fd)
				if path == 'socket:[' + inode + ']':
					try:
						f = file('/proc/' + pid + '/cmdline', 'rb')
						cmdline = f.read()
						f.close()
						cmdline = cmdline.split('\x00', 1)[0]
						cmdline = os.path.basename(cmdline)
					except:
						cmdline = None
					return pid, cmdline
		except:
			pass

	return None, None

if __name__ == '__main__':
	if len(sys.argv) < 2:
		print 'usage: %s <interface>' % (os.path.basename(sys.argv[0]), )
		sys.exit(0)

	p = pcap.pcapObject()
	#dev = pcap.lookupdev()
	dev = sys.argv[1]
	net, mask = pcap.lookupnet(dev)
	p.open_live(dev, 1600, 0, 100)
	p.setfilter('tcp', 0, 0)

	try:
		while 1:
			p.dispatch(1, print_packet)
	except KeyboardInterrupt:
		sys.excepthook(*sys.exc_info())
		print 'shutting down'
		print '%d packets received, %d packets dropped, %d packets dropped by interface' % p.stats()