#!/usr/bin/env python3

from dataclasses import dataclass
import argparse
from sys import argv, exit, stderr
from copy import copy
import pathlib
import formats

# Table from https://www.kernel.org/doc/html/latest/admin-guide/tainted-kernels.html:
_taint_table = [
	"0 G/P 1 proprietary module was loaded",
	"1 _/F 2 module was force loaded",
	"2 _/S 4 kernel running on an out of specification system",
	"3 _/R 8 module was force unloaded",
	"4 _/M 16 processor reported a Machine Check Exception (MCE)",
	"5 _/B 32 bad page referenced or some unexpected page flags",
	"6 _/U 64 taint requested by userspace application",
	"7 _/D 128 kernel died recently, i.e. there was an OOPS or BUG",
	"8 _/A 256 ACPI table overridden by user",
	"9 _/W 512 kernel issued warning",
	"10 _/C 1024 staging driver was loaded",
	"11 _/I 2048 workaround for bug in platform firmware applied",
	"12 _/O 4096 externally-built (“out-of-tree”) module was loaded",
	"13 _/E 8192 unsigned module was loaded",
	"14 _/L 16384 soft lockup occurred",
	"15 _/K 32768 kernel has been live patched",
	"16 _/X 65536 auxiliary taint, defined for and used by distros",
	"17 _/T 131072 kernel was built with the struct randomization plugin",
	]

# FIXME: Support for multi-bit taints? (i.e. P+O=4097 means likely weird driver…)
@dataclass
class TaintBit:
	position: int
	true_letter: str	# Taint set
	false_letter: str	# Taint not set
	description: str

	def __str__(self): return self.description + f' (bit {self.position})'


class TaintTable:
	def __init__(self, bits: dict[int, TaintBit]):
		self.bits = bits
		# Pre-compute the all-false string
		self.false_letters = ['' for _ in range(max(self.bits.keys())+1)]
		for bit in self.bits.values():
			self.false_letters[bit.position] = bit.false_letter
	
	def get_taints(self, value):
		# FIXME: might be slow?
		result = []
		l = value.bit_length()
		for pos,bit in self.bits.items():
			if pos >= l: continue
			if value & (1 << bit.position):
				result.append(bit)
		return sorted(result, key=lambda b:b.position)

	def get_taint_letters(self, taints):
		letters = copy(self.false_letters)
		for t in taints:
			letters[t.position] = t.true_letter
		return ''.join(letters)

def get_taint_mapping(table):
	result = {}
	for line in table:
		pos, letters, _dec, descr = line.split(' ', maxsplit=3)
		pos = int(pos)
		falselet, truelet = letters.split('/')
		if falselet == '_': falselet = ' '
		assert pos not in result
		result[pos] = TaintBit(position=pos, true_letter=truelet, false_letter=falselet, description=descr)
	return TaintTable(result)

TAINTS = get_taint_mapping(_taint_table)

def strict_check(taints, value) -> bool:
	correct_value = sum(1<<taint.position for taint in taints)
	return value == correct_value

def main():
	ap = argparse.ArgumentParser(description='Extracts human-friendly information about kernel taint',
			prog=argv[0],
			conflict_handler='resolve', # Needed to have -h be --human instead of help.
			)

	ap.usage = '\t%(prog)s [what] [format] [source]\n\t%(prog)s [--help]'
	ap.epilog = '%(prog)s terminates with status 1 if a taint was found, 0 if not. If strict interpretation was enabled, code 2 represents bad taint value.'
	ap.allow_abbrev = True
	
	# General options
	ap.add_argument('-S', '--strict', action='store_true',
			help='Interpret value strictly – fail if it is not a real taint value')

	what = ap.add_argument_group(title='What', description='Controls what information should be output')
	what.add_argument('-q', '--quiet', action='store_true',
			help='Suppres default output')
	what.add_argument('-d', '--descriptions', action='store_true',
			help='Show descriptions of values')
	what.add_argument('-l', '--letters', action='store_true',
			help='Show the letter string, as in kernel dump (default for script format)')
	what.add_argument('-b', '--bits', action='store_true',
			help='Show a list of bit orders, e.g. "1,4,12"')
	what.add_argument('-c', '--chktaint', action='store_true',
			help='Mimic kernel_chktaint output (default for human format)')

	how = ap.add_argument_group(title='Format', description='Format of the output')
	how.add_argument('-s', '--script', action='store_true',
			help='Output KEY=value, as in /etc/os-release')
	how.add_argument('-h', '--human', action='store_true',
			help='Use human-friendly format (default)')

	# This should be mutually exclusive group, but in order to have nice help, it is not.
	whence = ap.add_argument_group(title='Source', description='Where to get the value from (default: read /proc/sys/kernel/tainted)')
	whence.add_argument('-f', '--file', action='store', type=pathlib.Path, default=None,
			help='File to read the value from')
	whence.add_argument('-i', '--info', action='store_true',
			help='Only display all known taints')
	whence.add_argument('int', action='store', type=int, nargs='?', metavar='<int>',
			help='Use the specific value')
	whence.add_argument('-', action='store_true', dest='stdin',
			help='Read the value from stdin')

	args = ap.parse_args()
	# FIXME: Handle errors ourself for consistency (i.e. exit with 3, not 2!)

	# Resolve argument collisions
	def die(reason):
		ap.print_usage()
		print(reason, file=stderr)
		exit(3)
	def collision(reason:str, *params):
		if sum(params) >= 2: die(reason)

	collision('Can show either script- or human-friendly output, not both.',
			args.script,
			args.human)
	collision('Can either read taint value from file, stdin, argument or print whole table, not multiple.',
			args.file is not None,
			args.info,
			args.int is not None,
			args.stdin)
	if args.int is not None and args.int < 0: die('Only non-negative values are valid taint values.')
	if args.file is not None and not args.file.exists(): die('The specified file does not exist.')

	#############################
	#                           #
	#  Phew, arguments parsed!  #
	#                           #
	#############################

	if args.int is not None: taint_val = args.int
	elif args.file is not None:
		with open(args.file, 'r') as f:
			taint_val = int(f.readline())
	elif args.stdin: taint_val = int(input())
	elif args.info:
		taint_val = sum(1<<pos for pos in TAINTS.bits.keys())
	else:
		# No option specified, using /proc/sys/kernel/tainted
		with open('/proc/sys/kernel/tainted', 'r') as f:
			taint_val = int(f.readline())
	
	#### Value obtained! ####

	taints = TAINTS.get_taints(taint_val)

	# Script variable names and formatters
	# flag_name -> (script_name, formatter)
	# FIXME: Too much boilerplate

	output_data = {
		'descriptions': ('DESCRIPTIONS', formats.description_format),
		'bits': ('BITS', formats.bit_format),
		'letters': ('LETTERS', formats.letter_format),
		'chktaint': ('CHKTAINT', formats.chktaint_format),
		}

	default_contents = not any([args.quiet, args.descriptions, args.letters, args.bits, args.chktaint])
	default_format = not any([args.script, args.human])

	# Set defaults
	if default_format: args.human = True
	if default_contents and args.human: args.chktaint = True
	if default_contents and args.script: args.letters = True

	for cont in ['descriptions', 'bits', 'letters', 'chktaint']:
		if getattr(args, cont):
			output = []
			if args.script: output.append(output_data[cont][0]+'=')
			output.append(output_data[cont][1](taints, TAINTS))
			print(''.join(output))
	
	
	return_value = 1 if taints else 0
	if args.strict: return_value = 2 if not strict_check(taints, taint_val) else return_value
	# FIXME: Output when strict and not quiet?

	exit(return_value)





if __name__ == '__main__': main()