#!/usr/bin/env python
#
#  Wifiroamd: Automatically connect to wireless access points.

import os, string, re, glob, time, syslog, urllib, sys, signal, optparse


configuration = {
		'device' : 'auto',
		'scanretries' : 5,
		'scaninterval' : 0.7,
		'scriptdir' : '/etc/wifiroamd/scripts',
		'connectionsdir' : '/etc/wifiroamd/connections',
		'etcdir' : '/etc/wifiroamd',
		'seendir' : '/etc/wifiroamd/seen',
		'communicationsdir' : '/etc/wifiroamd/',
		}


#######################
class SigTermException(Exception): pass
class NoWirelessFound(Exception): pass
class RetryConnect(Exception): pass


#################
class ExceptHook:
	#################################################
	def __init__(self, useSyslog = 1, useStderr = 0):
		self.useSyslog = useSyslog
		self.useStderr = useStderr


	#######################################
	def __call__(self, etype, evalue, etb):
		import traceback, string
		tb = traceback.format_exception(*(etype, evalue, etb))
		tb = map(string.rstrip, tb)
		tb = string.join(tb, '\n')
		for line in string.split(tb, '\n'):
			if self.useSyslog:
				syslog.syslog(line)
			if self.useStderr:
				sys.stderr.write(line + '\n')


##################
class HelperClass:
	def __init__(self, configuration):
		self.configuration = configuration
		self.shutdownScript = None
		self.scriptEnv = {}
		self.logLevel = 0
		self.monitorBlacklist = []


	#############################
	def setLogLevel(self, level):
		'''Set the log level to the specified version.'''
		self.logLevel = level


	#####################
	def log(self, *args):
		'''Log a message, if 2 arguments, the first is the level.'''
		logLevel = 0
		if len(args) == 2:
			logLevel = args[0]
			args = ( args[1], )
		if logLevel > self.logLevel: return

		syslog.syslog(args[0])


	####################
	def getDevice(self):
		'''Figure out what wireless device to use.  If configuration specifies
		"auto", then a device is looked for in /sys/class/net.'''
		if self.configuration['device'] == 'auto':
			for sysDir in glob.glob('/sys/class/net/*/wireless'):
				dev = sysDir[15:]
				dev = string.split(dev, '/')[0]
				self.configuration['device'] = dev
				break
		if self.configuration['device'] == 'auto':
			fp = os.popen('/sbin/iwconfig 2>&1', 'r')
			for line in fp.readlines():
				m = re.match(r'^(\S+)\s+.*\s+ESSID:.*', line)
				if not m: continue
				self.configuration['device'] = m.group(1)
				break
			fp.close()
		return(self.configuration['device'])
	

	###############
	def scan(self):
		'''Scan for wireless devices, return a list of dictionaries about them.'''
		#  run scan
		fp = os.popen('iwlist "%s" scan 2>/dev/null' % self.getDevice(), 'r')
		data = {}
		scans = []
		for line in fp.readlines():
			self.log(6, 'iwlist line: "%s"' % repr(line))
			m = re.match(r'^\s+Cell (\d+) - Address: ([a-fA-F\d:]+)\s*$', line)
			if m:
				self.log(5, 'Matched Cell/Address line')
				if data:
					scans.append(data)
				data = {}
				data['Cell'] = m.group(1)
				data['Address'] = m.group(2)

			m = re.match(r'^\s+ESSID:"(.*)"\s*$', line)
			if m:
				self.log(5, 'Matched essid')
				data['ESSID'] = m.group(1)

			m = re.match(r'^\s+Mode:(.*\S)\s*$', line)
			if m:
				self.log(5, 'Matched mode')
				data['Mode'] = m.group(1)

			m = re.match(r'^\s+Signal level=(\S.*\S)\s*$', line)
			if m:
				self.log(5, 'Matched Signal level')
				data['Signal level'] = int(string.split(m.group(1))[0])

			m = re.match(r'^\s+Quality=([0-9/]+)\s+Signal level=(\S.*\S)\s*$',
					line)
			if m:
				self.log(5, 'Matched quality')
				lhs, rhs = map(float, string.split(m.group(1), '/'))
				data['Quality'] = lhs / rhs
				data['Signal level'] = int(string.split(m.group(2))[0])

			m = re.match(r'^\s+Encryption key:(on|off)\s*$', line)
			if m:
				self.log(5, 'Matched encryption')
				data['Encryption key'] = m.group(1)

			m = re.match(r'^\s+Bit Rate:(.*\S)\s*$', line)
			if m:
				self.log(5, 'Matched rate')
				data['Bit Rate'] = m.group(1)
		if data:
			scans.append(data)
		del(data)

		fp.close()

		return(scans)


	###################
	def runScans(self):
		'''Run the specified number of scans, returning a list of APs found,
		sorted by strength.'''
		#  wait for wireless device
		wirelessSysFile = os.path.join('/sys/class/net', self.getDevice())
		if not os.path.exists(wirelessSysFile):
			self.log(1, 'Wireless device is unavailable, waiting up '
					'to 30 seconds.')
			for i in xrange(30):
				time.sleep(1)
				if os.path.exists(wirelessSysFile): break
			if os.path.exists(wirelessSysFile):
				self.log(1, 'Found wireless device.')

		#  start scan
		self.log(3, 'Scanning for APs, %d times in %.1f seconds.'
				% ( self.configuration['scanretries'],
					self.configuration['scanretries']
						* self.configuration['scaninterval'] ))

		#  get a list of APs
		apDict = {}
		for i in xrange(self.configuration['scanretries'], 0, -1):
			for ap in self.scan():
				apDict[ap['Address']] = ap
			if i != 1:
				time.sleep(self.configuration['scaninterval'])

		#  generate sorted results
		foo = []
		for apInfo in apDict.values():
			#  skip APs with no ESSID
			if not apInfo.has_key('ESSID'):
				continue

			#  set up Quality value
			if apInfo.has_key('Quality'): quality = apInfo['Quality']
			else:
				if apInfo.has_key('Signal level'):
					quality = (apInfo['Signal level'] + 100) / 100.0
				else:
					self.log(2, 'AP with ESSID "%s" had no quality or signal'
							% apInfo['ESSID'])
					quality = 0

			#  add to list
			foo.append(( quality, apInfo ))
		foo.sort()
		foo.reverse()
		apList = map(lambda x: x[1], foo)

		#  log AP information
		if self.logLevel >= 3:
			self.log(3, 'Found the following AP information:')
			for ap in apList:
				self.log(3, 'AP: "%s"' % repr(ap))

		return(apList)


	#############################
	def findScript(self, apList):
		'''Look for a script for the APs, return None if not found.'''
		scriptDir = self.configuration['connectionsdir']
		for ap in apList:
			for script in (
					'mac:%s' % string.upper(ap['Address']),
					'mac:%s' % string.lower(ap['Address']),
					'essid:%s' % urllib.quote(ap['ESSID'], ''),
					):
				file = os.path.join(scriptDir, script)
				if os.path.exists(file):
					return(( file, ap ))
		return(( None, None ))


	############################################
	def runScript(self, script, arguments = ''):
		'''Run the specified script, or the scripts within the specified
		directory.  Return the environment values set by the script(s).'''
		if not os.path.exists(script):
			self.log(2, 'Skipping script "%s", no such file or directory.'
					% script)
			return({})

		#  log WIFIROAMD environment
		if self.logLevel >= 9:
			keys = os.environ.keys()
			keys.sort()
			for key in keys:
				if not key.startswith('WIFIROAMD_'): continue
				self.log(9, 'Start environment: "%s"="%s"'
						% ( key, repr(os.environ.get(key)) ))

		#  run directory of scripts
		if os.path.isdir(script):
			self.log(2, 'Running script directory "%s %s"'
					% ( script, arguments ))
			env = {}
			files = glob.glob(os.path.join(script, '*'))
			files.sort()
			for newScript in files:
				newEnv = self.runScript(newScript, arguments)
				env.update(newEnv)
				os.environ.update(newEnv)
			return(env)

		#  look for common package or save files.
		scriptBasename = os.path.basename(script)
		if (script.endswith('.dpkg-old') or 
				script.endswith('.old') or 
				script.endswith('.bak') or 
				script.endswith('.rpmsave') or 
				script.endswith('.rpmnew') or 
				script.endswith('~') or 
				scriptBasename.startswith('.') or 
				script.endswith('.rpmold')):
			self.log(3, 'Ignoring script because of extension: "%s"' % script)
			return({})

		#  run shell script
		self.log(3, 'Running script "%s %s"'
					% ( script, arguments ))
		stdin, stdout = os.popen4('bash -c "trap set EXIT; . \'%s\' %s 2>&1"'
					% ( script, arguments ))
		stdin.close()
		env = {}
		for line in stdout.readlines():
			if not line.startswith('WIFIROAMD_'): continue
			m = re.match(r'^(\S+)=(.*\S*)\s*$', line)
			if not m: continue

			value = m.group(2)
			if value and len(value) > 2 and value[0] == "'" and value[-1] == "'":
				#  strip out shell quoting
				value = value[1:-1]
				value = string.replace(value, "'\\''", "'")

			env[m.group(1)] = value
			self.log(8, 'Script environment: "%s"="%s"'
					% ( m.group(1), value ))
		stdout.close()

		return(env)


	########################################
	def monitorBlacklistPreen(self, apList):
		'''Remove monitor blacklist entries from the AP list.'''
		#  short circuit
		if not self.monitorBlacklist:
			self.log(6, 'Monitor blacklist is empty.')
			return(apList)

		origApList = apList
		apList = apList[:]
		for ap in self.monitorBlacklist:
			self.log(6, 'Preening monitor blacklist: "%s"' % repr(ap))
		for ap in self.monitorBlacklist:
			if not ap: continue
			for testAp in apList:
				if not testAp: continue
				if testAp.get('Address', 0) == ap.get('Address', 1):
					apList.remove(testAp)
					self.log(5, 'Monitor blacklist removed AP "%s"' % repr(ap))

		return(apList)


	##################
	def connect(self):
		'''Find an access-point and get connected to it.'''
		self.log(1, 'Starting connection')

		#  clear environment
		for key in os.environ.keys():
			if key.startswith('WIFIROAMD_') or key in (
					'IFACE', 'DEVICE', 'ESSID', 'MACADDRESS' ):
				del(os.environ[key])
		env = self.runScript(
				os.path.join(self.configuration['etcdir'], 'config-defaults'))
		os.environ.update(env)
		env = self.runScript(
				os.path.join(self.configuration['etcdir'], 'config'))
		os.environ.update(env)
		os.environ['WIFIROAMD_VERBOSE_LEVEL'] = str(self.logLevel)

		#  initialize device
		self.runScript(
				os.path.join(self.configuration['scriptdir'], 'reset.d'),
				'start')

		#  connect to access point
		apList = self.runScans()
		apList = self.monitorBlacklistPreen(apList)
		while 1:
			ap = None
			try:
				self.log(5, 'apList loop:')
				if not apList and self.monitorBlacklist:
					self.monitorBlacklist = []
					self.log(4, 'Resetting monitor blacklist.')
					raise RetryConnect
				for logAp in apList:
					self.log(5, '   "%s"' % repr(logAp))
				if not apList:
					self.log(3, 'apList is empty, breaking out of AP locate loop.')
					raise NoWirelessFound

				script, ap = self.findScript(apList)
				if script:
					self.log(3, 'Found script="%s" ap="%s"'
							% ( repr(script), repr(ap) ))
				else:
					self.log(2, 'No script found, trying first unencrypted AP')
					script = os.path.join(self.configuration['connectionsdir'],
							'default.d')

					#  connect to the first non-protected AP
					ap = None
					for checkAp in apList:
						if checkAp['Encryption key'] == 'off':
							self.log(1,
									'Selected non-encrypted AP "%(ESSID)s"/"%(Address)s"'
									% checkAp)
							ap = checkAp
							self.log(4, 'AP info: "%s"' % repr(ap))
							break

					#  create seen file
					if ap:
						self.log(3, 'seendir: "%s"'
								% repr(self.configuration.get('seendir')))
						self.log(3, 'Address: "%s"' % repr(ap.get('Address')))
						seenFile = os.path.join(self.configuration['seendir'],
								'mac:%s' % string.upper(ap['Address']) )
						if os.path.exists(seenFile):
							os.utime(seenFile, None)
						else:
							fp = open(seenFile, 'w')
							fp.write('#  ESSID: %s\n' % ap['ESSID'])
							fp.write('#  Seen on %s\n' % time.ctime().strip())
							fp.close()

				#  use "nowireless.d" script if no AP is found
				self.callUpScripts = 1
				if not ap: raise NoWirelessFound
			except NoWirelessFound:
				self.callUpScripts = 1
				self.log(1, 'No APs found, using "nowireless" script...')
				script = os.path.join(self.configuration['connectionsdir'],
						'nowireless.d')
				apList = []
				self.ap = None

			#  save off script for shutdown use
			self.shutdownScript = script

			#  set up environment
			os.environ['IFACE'] = self.configuration['device']
			os.environ['DEVICE'] = self.configuration['device']
			os.environ['WIFIROAMD_DEVICE'] = self.configuration['device']
			if ap:
				apEssid = ap.get('ESSID', '')
				os.environ['ESSID'] = apEssid
				os.environ['WIFIROAMD_ESSID'] = apEssid
				os.environ['ESSID_ESCAPED'] = urllib.quote(apEssid, '')
				os.environ['WIFIROAMD_ESSID_ESCAPED'] = urllib.quote(apEssid, '')
				os.environ['MACADDRESS'] = ap.get('Address', '')
				os.environ['WIFIROAMD_MACADDRESS'] = ap.get('Address', '')

			#  run start scripts
			scriptList = [
					( script, 'start' ),
					]
			if self.callUpScripts:
				scriptList.append(
						( os.path.join(self.configuration['scriptdir'], 'up.d'),
							'start' ))
				scriptList.append(( script, 'up' ))
			self.scriptEnv = {}
			doBlacklist = 0
			while scriptList:
				script, argument = scriptList[0]
				del(scriptList[0])
				env = self.runScript(script, argument)

				#  log environment
				for envKey, envValue in env.items():
					self.log(9, 'Env: "%s" = "%s"' % ( envKey, envValue ))

				#  add a next script to the list
				if env.get('WIFIROAMD_NEXTSCRIPT'):
					scriptList.insert(0, ( env.get('WIFIROAMD_NEXTSCRIPT'),
							'start' ))
					del(env['WIFIROAMD_NEXTSCRIPT'])

				#  blacklist AP
				if env.get('WIFIROAMD_BLACKLIST') == '1':
					doBlacklist = 1
					break

				#  update environment
				self.scriptEnv.update(env)
				os.environ.update(env)

			#  do not connect to this AP
			if doBlacklist:
				self.log(2, 'This AP is blacklisted, trying another...')
				apList.remove(ap)
				self.ap = None
				continue

			#  connected to AP
			self.ap = ap
			break


	#####################
	def disconnect(self):
		'''After we have lost connectivity, run the script with "stop".'''
		self.log(1, 'Shutting down connection')

		#  run stop scripts
		scriptList = [
				self.shutdownScript,
				]
		if self.callUpScripts:
			scriptList.append(os.path.join(self.configuration['scriptdir'],
					'down.d'))
		while scriptList:
			script = scriptList[0]
			del(scriptList[0])
			if not script: continue
			env = self.runScript(script, 'stop')

			#  add a next script to the list
			if env.get('WIFIROAMD_NEXTSCRIPT'):
				scriptList.insert(0, env.get('WIFIROAMD_NEXTSCRIPT'))
				del(env['WIFIROAMD_NEXTSCRIPT'])

			#  update environment
			self.scriptEnv.update(env)
			os.environ.update(env)

		self.scriptEnv = {}
		self.shutdownScript = None


	##################
	def monitor(self):
		'''Run a monitor script, return 1 when one of them reports an
		exit status of non-0.'''
		self.log(2, 'Monitoring')

		#  run stop scripts
		if os.environ.get('WIFIROAMD_MONITOR_PINGGW') == '1':
			self.log(3, 'Monitoring: Pinging Gateway')
			scriptList = [
					os.path.join(self.configuration['scriptdir'], 'monitor_pinggw.d')
					]
			if self.shutdownScript: scriptList.insert(0, self.shutdownScript)
			self.log(5, 'Script list: "%s"' % repr(scriptList))
			while scriptList:
				script = scriptList[0]
				del(scriptList[0])
				env = self.runScript(script, 'monitor_pinggw')

				#  add a next script to the list
				if env.get('WIFIROAMD_NEXTSCRIPT'):
					scriptList.insert(0, env.get('WIFIROAMD_NEXTSCRIPT'))
					del(env['WIFIROAMD_NEXTSCRIPT'])

				#  update environment
				self.scriptEnv.update(env)
				os.environ.update(env)

				if env.get('WIFIROAMD_SCRIPTEXITCODE', '0') != '0':
					self.log(4,
							'Monitor script "%s" returned "%s", closing connection.'
							% ( script, env.get('WIFIROAMD_SCRIPTEXITCODE') ))
					return(False)

		return(True)


	########################
	def sigpidwrapper(self):
		#  check pid-file
		pidFile = '/var/run/wifiroamd.pid'
		if os.path.exists(pidFile):
			fp = open(pidFile, 'r')
			line = fp.readline()
			fp.close()
			if line.strip():
				pid = int(line)
				try:
					os.kill(pid, 0)
					s = 'Another process, id=%s, already exists.  Aborting.' % pid
					sys.stderr.write(s + '\n')
					syslog.syslog(s)
					sys.exit(1)
				except OSError:
					os.remove(pidFile)

		#  create pid-file
		fp = open(pidFile, 'w')
		fp.write('%s\n' % os.getpid())
		fp.close()

		#  call main with signal handler
		signal.signal(signal.SIGTERM, self.sigterm)
		try:
			self.main()
		except SigTermException:
			signal.signal(signal.SIGTERM, signal.SIG_DFL)
			self.log(1, 'Got TERM signal, cleaning up.')
			helper.disconnect()

		#  remove pid-file
		os.remove(pidFile)


	#################################
	def sigterm(self, signum, frame):
		'''Signal handler for TERM.'''
		raise SigTermException


	###############
	def main(self):
		'''Main code.'''
		#  only run as root
		if os.getuid() != 0:
			s = 'Must run as root, aborting.'
			sys.stderr.write(s + '\n')
			syslog.syslog(s)
			sys.exit(1)

		disabledFile = os.path.join(self.configuration['communicationsdir'],
				'disabled')
		while 1:
			#  wait for disabled file
			if os.path.exists(disabledFile):
				self.log(1, 'Waiting for disabled file to go away.')
				while os.path.exists(disabledFile):
					time.sleep(1)

			#  loop trying connections
			while 1:
				try:
					helper.connect()
					break
				except RetryConnect: pass

			#  monitor
			monitorSuccesses = 0
			while 1:
				if not helper.monitor():
					break
				if os.path.exists(disabledFile):
					self.log(2, 'Disabled file found, disconnecting.')
					break

				sleepSecs = float(os.environ.get('WIFIROAMD_MONITOR_INTERVAL', 10))
				helper.log(3, 'Sleeping %.1f before next monitor' % sleepSecs)
				monitorSuccesses = monitorSuccesses + 1
				time.sleep(sleepSecs)

			#  monitor blacklist
			if monitorSuccesses > 10:
				helper.monitorBlacklist = []
				helper.log(4, 'Monitor succeeded, clearing blacklist')
			else:
				helper.monitorBlacklist.append(helper.ap)
				helper.log(4, 'Monitor failed, adding to blacklist: "%s"'
						% repr(helper.monitorBlacklist))

			helper.disconnect()


	####################
	def parseOpts(self):
		parser = optparse.OptionParser()
		parser.add_option('-v', '--verbose', action = 'count', dest = 'verbosity')
		options, args = parser.parse_args()
		self.setLogLevel(options.verbosity)


##############################
def checkForHotplugConflict():
	if glob.glob('/etc/sysconfig/network-scripts/ifcfg-eth*'):
		syslog.syslog('WARNING: Presence of '
				'"/etc/sysconfig/network-scripts/ifcfg-*" scripts for the '
				'wireless and wired interfaces may cause conflicts and weirdness '
				'with wifiroamd.  It is recommended that you rename these files '
				'(do NOT rename ifcfg-lo) to "disabled-ifcfg-$IFNAME".')


############
#  main code
syslog.openlog('wifiroamd')
sys.excepthook = ExceptHook(useSyslog = 1, useStderr = 1)
checkForHotplugConflict()
helper = HelperClass(configuration)
helper.parseOpts()
helper.sigpidwrapper()
sys.exit(0)
