Set names of threads for better debugging.

[oddments/cappie.git] / cappie / __init__.py

#!/usr/bin/python
###############################################################################
#                                                                             #
# Cappie                                                                      #
# Copyright (C) 2010 Michael Tremer                                           #
#                                                                             #
# This program 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, either version 3 of the License, or           #
# (at your option) any later version.                                         #
#                                                                             #
# This program 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.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

import logging
import logging.handlers
import pcapy
import time

from ConfigParser import ConfigParser
from threading import Thread

import protocol
import queue
import util

from errors import *

def getAllInterfaces():
	filters = ("lo", "any")
	ret = []
	for dev in pcapy.findalldevs():
		if not dev in filters:
			ret.append(dev)
	return ret

class Cappie(object):
	def __init__(self):
		self.__interfaces = []

		self.log = logging.getLogger("cappie")
		self.log.setLevel(logging.INFO)

		# Log to console
		handler = logging.StreamHandler()
		handler.setFormatter(logging.Formatter("%(levelname)7s %(message)s"))
		self.log.addHandler(handler)

		# Setup syslog
		handler = logging.handlers.SysLogHandler("/dev/log")
		handler.setFormatter(logging.Formatter("cappie: %(message)s"))
		self.log.addHandler(handler)

		self.queue = queue.Queue(self.log)

		self.log.info("Cappie successfully started")

	def __del__(self):
		self.shutdown()
		self.log.info("Exiting")

	def setDebug(self, debug):
		if debug:
			self.log.setLevel(logging.DEBUG)
		else:
			self.log.setLevel(logging.INFO)

	def addInterface(self, dev, **kwargs):
		if not dev in getAllInterfaces():
			raise InterfaceError, "No such interface %s" % dev

		kwargs["cappie"] = self

		iface = Interface(dev, **kwargs)
		self.__interfaces.append(iface)

	def run(self):
		if not self.__interfaces:
			raise RuntimeError, "No interfaces were configured"

		# Start queue
		self.queue.start()

		# Start a thread for each interface
		for iface in self.__interfaces:
			iface.start()

		while True:
			if not self.queue.is_alive():
				self.log.critical("Queue thread died unexpectedly.")
				return

			for iface in self.__interfaces:
				if not iface.is_alive():
					self.log.critical("Thread died unexpectedly. %s" % iface.dev)
					return
				time.sleep(60)

	def readConfig(self, configfile):
		config = ConfigParser()
		config.read([configfile])

		global_opts = {}
		if config.has_section("global"):
			for option, value in config.items("global"):
				global_opts[option] = value

			config.remove_section("global")

		for iface in config.sections():
			options = {}
			for option, value in config.items(iface):
				options[option] = value
			self.addInterface(iface, **options)

	def shutdown(self):
		for iface in self.__interfaces:
			iface.shutdown()

		self.queue.shutdown()


class Interface(Thread):
	heartbeat = 0.1

	def __init__(self, dev, cappie, promisc=False, mtu=1500):
		Thread.__init__(self)

		self.cappie = cappie
		self.dev = dev
		self.log = self.cappie.log
		self.mtu = mtu
		self.promisc = promisc
		self.queue = self.cappie.queue

		self.db = Database(self)

		self.log.debug("Created new interface %s" % self.dev)
		
		self.__running = True

	def _callback(self, header, data):
		self.log.debug("Received packet on %s" % self.dev)
		try:
			p = protocol.decode_packet(data)
		except PacketTypeError, e:
			self.log.error("Got unknown packet: %s" % e)
			return
		except DecodeError, e:
			self.log.warning("Got decoding error: %s" % e)
			return

		# Dump packet information
		for key, val in p.items():
			self.log.debug("  %s: %s" % (key, val))

		if not self.db.has(p["source_address"]):
			self.db.put(p["source_address"], "SOURCE_IP_ADDRESS", p["source_ip_address"])

	def run(self):
		self.log.info("Starting interface %s" % self.dev)

		util.setprocname("interface %s" % self.dev)

		self.db.open()

		p = pcapy.open_live(self.dev, self.mtu, self.promisc, 0)
		p.setfilter(self.filter)
		#p.loop(0, self._callback)

		p.setnonblock(1)
		while True:
			if not self.__running:
				self.db.close()
				return
			
			if p.dispatch(1, self._callback):
				continue

			time.sleep(self.heartbeat)

	def shutdown(self):
		if not self.__running:
			return

		self.log.debug("Sending shutdown signal to %s" % self.dev)
		self.__running = False

	@property
	def filter(self):
		return "arp or rarp"


class Database(object):
	def __init__(self, interface):
		self.interface = interface
		self.dev = self.interface.dev
		self.log = self.interface.log

		self.__data = {}

	def open(self):
		self.log.debug("Opened database for %s" % self.dev)

	def close(self):
		self.log.debug("Closing database for %s" % self.dev)
		print self.__data

	def get(self, mac):
		if self.has(mac):
			return self.__data[mac]

	def has(self, mac):
		return self.__data.has_key(mac)

	def put(self, mac, key, val):
		if not self.has(mac):
			self.__data[mac] = {}

		# TODO Check key for sanity

		self.__data[mac][key] = val