#!/usr/bin/env python3
#
# k8spodwatchd.py - Track pod state changes and emit alert events as required (from within a cluster)
#
# Copyright 2022 Todd Shadburn <tshadburn@opencomputingsolutions.com>
#
# 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 2 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, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#


import os
import sys
import time
import datetime
import json
import yaml
import logging
import netinfo
from kubernetes import client, config, watch
from kubernetes.client import Configuration


# Interval between checks
CHECK_INTERVAL_SECONDS = 60
 

# Emit alert if pod remains in a 'bad' state for longer than
# this amount of time
DEFAULT_ALERT_DURATION_SECONDS = 115


# Pods that are in one of these phases will not event
# (considered to be in a 'good' state)
nonalert_phases = [
	'Running',
	'Succeeded',
]


# Class to create a client object from a specific kubeconfig file
class KubeConfig(object):
	def __init__(self, configuration_yaml):
		self.configuration_yaml = configuration_yaml
		self._configuration_yaml = None

	@property
	def config(self):
		with open(self.configuration_yaml, 'r') as f:
			if self._configuration_yaml is None:
				self._configuration_yaml = yaml.safe_load(f)
		return self._configuration_yaml

	@property
	def client(self):
		k8_loader = config.kube_config.KubeConfigLoader(self.config)
		call_config = type.__call__(Configuration)
		k8_loader.load_and_set(call_config)
		Configuration.set_default(call_config)
		return client.CoreV1Api()


# A state tracking class
class StateTracker(object):
	def __init__(self, **kwargs):
		self.state = {}
		self.state_change_callback = None
		if 'state_change_callback' in kwargs:
			self.state_change_callback = kwargs['state_change_callback']
		return

	def exists(self, objid=None):
		if objid in self.state:
			return True
		return False

	def current_duration(self, objid=None):
		if not objid in self.state:
			raise ValueError('That object id is not being tracked')
		return (datetime.datetime.now() - self.state[objid]['track_start'])

	def set_data(self, objid=None, data=None):
		if not objid in self.state:
			raise ValueError('That object id is not being tracked')
		self.state[objid]['data'] = data

	def get_data(self, objid=None):
		if not objid in self.state:
			raise ValueError('That object id is not being tracked')
		return self.state[objid]['data']

	def add(self, objid=None, state=None, **kwargs):
		if objid in self.state:
			raise ValueError('That object id is already being tracked')
		d = {
			'id': objid,
			'track_start': datetime.datetime.now(),
			'last_change': datetime.datetime.now(),
			'state': state,
			'data': None,
		}
		if 'data' in kwargs:
			d['data'] = kwargs['data']
		self.state[objid] = d
		return

	def update(self, objid=None, state=None, **kwargs):
		if not objid in self.state:
			raise ValueError('That object id is not being tracked')
		self.state[objid]['state'] = state
		self.state[objid]['last_change'] = datetime.datetime.now()
		if 'data' in kwargs:
			self.state[objid]['data'] = kwargs['data']
		return

	def remove(self, objid=None, **kwargs):
		if not objid in self.state:
			raise ValueError('That object id is not being tracked')
		del self.state[objid]
		return



def pod_container_ready_counts(pod=None):
	total = 0
	ready = 0
	if pod is None or pod.status.container_statuses is None:
		return [ready,total]
	for c in pod.status.container_statuses:
		if c.started == True:
			total += 1
		if c.ready:
			ready += 1
	return [ready,total]


def pod_container_state_counts(pod=None):
	total = 0
	running = 0
	termed = 0
	waiting = 0
	for c in pod.status.container_statuses:
		if c.started == True:
			total += 1
		if c.state.running:
			running += 1
		if c.state.terminated:
			termed += 1
		if c.state.waiting:
			waiting += 1
	return [total,running,termed,waiting]


def debug_pod(cluster=None, pod=None):
	print("DEBUG",','.join([
		cluster,
		pod.metadata.namespace,
		pod.metadata.name,
		pod.metadata.creation_timestamp.isoformat(),
		pod.status.phase,
		'/'.join('{0}'.format(n) for n in pod_container_ready_counts(pod)),
	]))
	return


def pod_get_controller_info(pod=None):
	if not pod:
		return (None,None)
	if hasattr(pod.metadata, 'owner_references'):
		for owner in pod.metadata.owner_references:
			if owner.controller == True:
				return (owner.kind,owner.name)
	return (None,None)


# Here we go
logger = logging.getLogger('k8spodwatchd.py')
default_log_args = {
	'level': logging.DEBUG if os.environ.get('DEBUG', False) else logging.INFO,
	'format': '%(asctime)s [%(levelname)s] %(name)s - %(message)s',
	'datefmt': '%Y-%m-%dT%H:%M:%SZ',
	#'force': True,	# only for Python 3.8+
}
logging.basicConfig(**default_log_args)


cluster = os.environ.get('CONFIG_K8S_CLUSTER_NAME', 'unknown')
ne_hostname = os.environ.get('CONFIG_NETINFO_HOSTNAME', 'localhost')
option_enable_initial_alert = os.environ.get('CONFIG_ENABLE_INITIAL_ALERT', False)
option_alert_duration_seconds = os.environ.get(
					'CONFIG_ALERT_DURATION_SECONDS',
					DEFAULT_ALERT_DURATION_SECONDS
)

nc_client = None
try:
	nc_client = netinfo.NetinfoClient(
		host=ne_hostname,
		cert='netinfo.cert',
		key='netinfo.key'
	)
except:
	logger.warning('failed initial connection to event sink, will reconnect later')
	pass

last_check = datetime.datetime.now(datetime.timezone.utc)
loop_start_time = last_check
tracker = StateTracker()
while True:
	st = time.time()

	# Perform checks
	evlist = []

	# Use a config file
	#cfg = KubeConfig(configuration_yaml='/home/someone/.kube/kubeconfig-%s' %(cluster))
	#api = cfg.client

	# Use the in-cluster config (relies on ServiceAccount and RBAC Role)
	config.load_incluster_config()
	api = client.CoreV1Api()

	res = api.list_pod_for_all_namespaces(watch=False)
	for pod in res.items:
		#debug_pod(cluster=cluster, pod=pod)
		objid = ','.join([cluster,pod.metadata.namespace,pod.metadata.name])
		ckind,cname = pod_get_controller_info(pod=pod)
		if ckind and ckind == 'Job':
			if hasattr(pod, 'start_time'):
				if pod.start_time < loop_start_time:
					logger.info('Skipping job pod %s as it is older than our startup time' % (objid))
					continue
		c_ready,c_total = pod_container_ready_counts(pod)
		if pod.status.phase not in nonalert_phases or c_ready != c_total:
			# pod is unhealthy somehow
			if not tracker.exists(objid):
				# Pod changed state
				tracker.add(objid, state=pod.status.phase)
				logger.info('%s failed (%s)' % (objid,pod.status.phase))
				if option_enable_initial_alert:
					msg = 'Pod %s in the %s namespace is in an unhealthy state(%s). Please check this issue.' % (
						pod.metadata.name,
						pod.metadata.namespace,
						pod.status.phase,
					)
					evlist.append(netinfo.NetinfoEvent(
						hostname=cluster,
						object=pod.metadata.name,
						flags=1,
						rc=1,
						data=str(msg),
					))
					logger.info(str(msg))
			else:
				# Pod still in "bad" state
				tracker.update(objid, state=pod.status.phase)
				#logger.warning('%s still bad (%s)' % (objid,pod.status.phase))
				if tracker.current_duration(objid) > datetime.timedelta(seconds=option_alert_duration_seconds):
					d = tracker.get_data(objid)
					#print('%s d=%s' % (objid,str(d)))
					if not d or d['alerting'] != True:
						tracker.set_data(
							objid,
							data={
								'alerting': True,
							}
						)
						msg = 'Pod %s in the %s namespace is in an unhealthy state (phase=%s ready=%d/%d). Please check this issue.' % (
							pod.metadata.name,
							pod.metadata.namespace,
							pod.status.phase,
							c_ready,
							c_total
						)
						evlist.append(netinfo.NetinfoEvent(
							hostname=cluster,
							object=pod.metadata.name,
							flags=1,
							rc=2,
							data=str(msg),
						))
						logger.info(str(msg))
		else:
			# pod is health
			if tracker.exists(objid):
				# Pod changed to a "good" state
				logger.info('%s recovered (%s)' % (objid,pod.status.phase))
				msg = 'Pod %s in the %s namespace has returned to a healthy state (phase=%s ready=%d/%d).' % (
					pod.metadata.name,
					pod.metadata.namespace,
					pod.status.phase,
					c_ready,
					c_total
				)
				evlist.append(netinfo.NetinfoEvent(
					hostname=cluster,
					object=pod.metadata.name,
					flags=1,
					rc=0,
					data=str(msg),
				))
				logger.info(str(msg))
				tracker.remove(objid)

	# Send any events that have been accumulated
	if len(evlist):
		for ev in evlist:
			try:
				nc_client.send_event(ev)
				logger.debug(ev)
			except:
				logger.warning('failed to send event, reconnecting')
				try:
					nc_client.close()
				except:
					# close as a best-effort
					pass
				nc_client = netinfo.NetinfoClient(
					host=ne_hostname,
					cert='netinfod.cert',
					key='netinfod.key'
				)
				nc_client.send_event(ev)
				logger.debug(ev)
				
		evlist = []
	# cleanup
	del api
	last_check = datetime.datetime.now(datetime.timezone.utc)

	# delay
	et = time.time()
	delay = CHECK_INTERVAL_SECONDS - (et-st)
	if delay < 0:
		logger.warning('poll operation is taking longer that the check interval delay=%0.3f' %(delay))
		delay = 0
	time.sleep(delay)

sys.exit(0)