#! /usr/bin/python3

import sys
import os

# Update sys.path (PYTHONPATH) to include the lib directory.
# Usually this script will be in $HOME/bin, and libtiny.py
# is in $HOME/lib. Typically $HOME/bin already is in PATH.
# $HOME/lib is not always in PYTHONPATH.
bindir = os.path.join( os.path.dirname(__file__) )
topdir = os.path.realpath( os.path.join( bindir, '..') )
libdir = os.path.join( topdir, 'lib' )
if libdir not in sys.path:
	sys.path.insert(0,libdir)

# Version in libdir takes precedence over system version
from libarkmon					import DB
from libarkmon.libtiny			import SAY
from libarkmon.libarkmon_helper	import attr_data

#+
# NAME:
#	arkmon-fix.py
# PURPOSE:
#	Maintenance of arkmon database
# CALLING SEQUENCE:
# OPTIONAL INPUTS:
#	--merge-new		merge nodes (used to fix node name changes
#					that are not automatically picked up by
#					a boot packet).
# RESTRICTIONS:
# PROCEDURE:
# MODIFICATION HISTORY:
#	MAY-2025, Paul Hick (UCSD/CAIDA)
#-

if __name__ == '__main__':

	from optparse import OptionParser

	version = 'V3.00.000 -- Paul Hick (UCSD/CAIDA; pphick@caida.org) -- 08-Mar-2025'
	usage = "%prog [--db DB ...] NODE1 NODE2 ..."

	parser = OptionParser(usage=usage,version=version)

	default_db   = 'arkmon-dev'
	default_user = 'arkmon-dev'
	default_attr_names = sorted( list( attr_data() ) )

	parser.add_option('-v', '--verbose'	,
		dest		= 'verbose'			,
		action		= 'count'			,
		default		= 0					,
		help		= 'verbose output'	,
	)

	parser.add_option('-n', '--dry-run'	,
		dest		= 'dryrun'			,
		action		= 'store_true'		,
		default		= False				,
		help		= 'make dryrun'		,
	)

	parser.add_option('', '--db'		,
		dest		= 'db'				,
		action		= 'store'			,
		default		= default_db		,
		help		= 'db name on ghul.caida.org, default: %s'%default_db,
	)

	parser.add_option('', '--?db'		,
		dest		= 'what_db'			,
		action		= 'store_true'		,
		default		= False				,
		help		= 'print postgres db information',
	)

	parser.add_option('', '--user'		,
		dest		= 'user'			,
		action		= 'store'			,
		default		= default_user		,
		help		= 'db user on ghul.caida.org, default: %s'%default_user,
	)

	parser.add_option('--show-pkt'		,
		dest        = 'show_pkt'		,
		type        = 'int'				,
		action      = 'store'			,
		default     = None				,
		help        = 'show pkt with specified timestamp',
	)

	parser.add_option('--delete-pkt'	,
		dest        = 'delete_pkt'		,
		type        = 'int'				,
		action      = 'store'			,
		default     = None				,
		help        = 'delete pkt with specified timestamp',
	)

	parser.add_option('', '--merge-new'	,
		dest		= 'merge_new'		,
		action		= 'store'			,
		default		= None				,
		help		= 'name of new node to be merged with existing old node',
	)

	parser.add_option('', '--hwmac-dups',
		dest		= 'hwmac_dups'		,
		action		= 'store_true'		,
		default		= False				,
		help		= 'search for hwmac dups',
	)

	options, args = parser.parse_args()

	say = SAY(
		label   = 'arkmon-fix'		,
		verbose = options.verbose	,
		dryrun  = options.dryrun                                                                        ,
	)

	db = DB( user=options.user, name=options.db )
	if not db.open:
		say.die( 'failed to open db "%s"'%options.db )
	say.yell( db.message )
	if options.what_db:
		sys.exit()

	if not options.hwmac_dups:
		if len(args) == 0:
			say.die( 'specify node name' )
		node = args[0]

		if not db.node_exists( node ):
			say.die( '%s does not exist'%node )

		if db.node_has_state( node ):
			say.yell( '%s has state'%node )

		if db.node_has_history( node ):
			say.yell( '%s has history'%node )

	if options.show_pkt:
		pkt = db.get_pkt( node, options.show_pkt )
		for attr in sorted( list(pkt) ):
			print ( '%-20s: %s'%(attr, pkt[attr]) )

	elif options.delete_pkt:
		db.delete_pkt( node, options.delete_pkt )

	elif options.merge_new:
		# Merging is done if a new node appears, that should be appended
		# to an existing node. The history file for the new file probably
		# is small, so we add the new history file to the old one; then
		# change the filename using the new node name.

		old_node = node

		if not options.merge_new:
			say.die( 'specify new node to be merged with old node %s'%old_node )
		new_node = options.merge_new

		if not db.node_exists( new_node ):
			say.die( '%s does not exist'%new_node )

		if db.node_has_state( new_node ):
			say.yell( '%s has state'%new_node )

		if db.node_has_history( new_node ):
			say.yell( '%s has history'%new_node )

		# Get HWmac (should be the same for both nodes)

		vv = db.get_attr_for_node( old_node, 'hwmac' )
		ww = db.get_attr_for_node( new_node, 'hwmac' )

		if vv != ww:
			say.die( 'different MAC: %s: %s  %s: %s'%(old_node, vv, new_node, ww) )

		say.say( 'merge new node "%s" with old node "%s" (mac=%s)'%(new_node, old_node, vv) )

		# Pick up all data for new_node
		data = { attr: db.dump(new_node, attr, sort='asc') for attr in db.attr_in_db }

		tt_old = db.get_attr_for_node( old_node, 'pkttype', get_time=True )
		tt_new = data['pkttype'][0][0]
		say.yell( '%s stops  at %s'%(old_node, tt_old) )
		say.yell( '%s starts at %s'%(new_node, tt_new) )
		if tt_new <= tt_old:
			say.die( 'data for two nodes overlap' )

		# Old node might be in graveyard
		if not db.node_has_state( old_node ):
			db.restore_node_state( old_node )

		# Add data for new node to old node
		for attr in db.attr_in_db:
			for row in data[attr]:
				db.insert( old_node, attr, row[1], row[0] )

		# Delete data for new node (they have been added to old_node)
		db.clear_node_state( new_node )
		for attr in db.attr_in_db:
			db.clear_node_history( new_node, attr)

		# new_node has now been deleted from the db, so
		# old_node can be changed to new_node
		db.change_hostname( { 'node': new_node, 'hostname': old_node } )

	elif options.hwmac_dups:
		# dict 'zrh3-ch': (1561142729, 'b8:27:eb:88:a9:2f')
		state = db.get_state_for_attr( 'hwmac', incl_none=False, get_pair=True )

		# list ('mce-us', 1747267119, '16:dc:6c:7b:21:95')
		history = db.read( 'SELECT node, time, value FROM hwmac', all=True )

		# All hwmac
		hwmacs = { state[node][1] for node in state } | { row[2] for row in history }

		data = {}

		for node in state:
			hwmac = state[node][1]
			if hwmac not in data:
				data[hwmac] = []
			data[hwmac].append(node)

		for row in history:
			hwmac = row[2]
			if hwmac not in data:
				data[hwmac] = []
			data[hwmac].append(row[0])

		for hwmac in data:
			if len(data[hwmac]) > 1:
				print ( hwmac+' '+' '.join(data[hwmac]) )
		
	else:
		say.yell('specify --attr ATTR, --time TIME, --dups, --sync, --merge')

	db.close()

	sys.exit(0)
