#!/usr/bin/python3

# MooseFS_License

# mfslicinfo - display MooseFS licence info from one or more master servers.
#
# Usage:
#   mfslicinfo [-h] [-v] host[:port] [host[:port] ...]
#
# Each positional argument is queried independently and its licence is shown.
# A failure on one host does not stop the others.

import os
import sys

if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 4):
	print("Unsupported python version, minimum required version is 3.4")
	sys.exit(1)

import socket
import struct
import getopt
import json
import time

### from common.constants_ac import * # MFS_INLINE_IMPORT - inlined below:
### Imported inline from common/constants_ac.py:
VERSION = "5.1.0"

PROTO_BASE = 0

DEFAULT_MASTERNAME          = "mfsmaster"
DEFAULT_MASTER_CLIENT_PORT  = 9421
DEFAULT_MASTER_CONTROL_PORT = 9419

BIN_PATH = "/usr/bin"

### End of inline import from common/constants_ac.py



ANTOAN_NOP = 0
CLTOMA_INFO = PROTO_BASE + 510
MATOCL_INFO = PROTO_BASE + 511
CLTOMA_LICENCE_DATA = PROTO_BASE + 1000
MATOCL_LICENCE_DATA = PROTO_BASE + 1001

LICVER_CE = -1


class MFSCommunicationError(RuntimeError):
	pass


def humanize_number(number, sep=''):
	number *= 100
	scale = 0
	while number >= 99950:
		number = number // 1024
		scale += 1
	if number < 995 and scale > 0:
		b = (number + 5) // 10
		nstr = "%u.%u" % divmod(b, 10)
	else:
		b = (number + 50) // 100
		nstr = "%u" % b
	if scale > 0:
		return "%s%s%siB" % (nstr, sep, "-KMGTPEZY"[scale])
	return "%s%sB" % (nstr, sep)


def get_string_from_packet(data, pos, err):
	rstr = b""
	shift = 0
	if 1 + pos > len(data):
		err = 1
	if err == 0:
		shift = data[pos]
		if shift + 1 + pos > len(data):
			err = 1
		else:
			rstr = data[pos + 1:pos + 1 + shift]
	return (rstr.decode('utf-8', 'replace'), pos + shift + 1, err)


def get_longstring_from_packet(data, pos, err):
	rstr = b""
	shift = 0
	if 2 + pos > len(data):
		err = 1
	if err == 0:
		shift = data[pos] * 256 + data[pos + 1]
		if shift + 2 + pos > len(data):
			err = 1
		else:
			rstr = data[pos + 2:pos + 2 + shift]
	return (rstr.decode('utf-8', 'replace'), pos + shift + 2, err)


def version_convert(versionxyzp):
	if versionxyzp >= (4, 0, 0) and versionxyzp < (4, 11, 0):
		return (versionxyzp, 0)
	if versionxyzp >= (2, 0, 0):
		return ((versionxyzp[0], versionxyzp[1], versionxyzp[2] // 2), versionxyzp[2] & 1)
	if versionxyzp >= (1, 7, 0):
		return (versionxyzp, 1)
	if versionxyzp > (0, 0, 0):
		return (versionxyzp, 0)
	return (versionxyzp, -1)


def version_str(versionxyzp):
	xyz, pro = version_convert(versionxyzp)
	s = "%u.%u.%u" % xyz
	if pro == 1:
		s += " PRO"
	if s == "0.0.0":
		s = "unknown"
	return s, pro


class MFSConn:
	def __init__(self, host, port, timeout=2):
		self.host = host
		self.port = port
		self.timeout = timeout
		self.socket = None
		self._connect()

	def _connect(self):
		self.socket = socket.socket()
		self.socket.settimeout(self.timeout)
		self.socket.connect((self.host, self.port))

	def close(self):
		if self.socket:
			try:
				self.socket.close()
			except Exception:
				pass
			self.socket = None

	def __del__(self):
		self.close()

	def _send(self, msg):
		total = 0
		while total < len(msg):
			sent = self.socket.send(msg[total:])
			if sent == 0:
				raise MFSCommunicationError("socket connection broken")
			total += sent

	def _recv(self, leng):
		buf = b""
		while len(buf) < leng:
			chunk = self.socket.recv(leng - len(buf))
			if len(chunk) == 0:
				raise MFSCommunicationError("socket connection broken")
			buf += chunk
		return buf

	def command(self, cmdout, cmdin, dataout=None):
		payload = dataout if dataout else b""
		msg = struct.pack(">LL", cmdout, len(payload)) + payload
		self._send(msg)
		while True:
			header = self._recv(8)
			cmd, length = struct.unpack(">LL", header)
			if cmd == cmdin:
				return self._recv(length), length
			if cmd != ANTOAN_NOP:
				raise MFSCommunicationError(
					"unexpected answer: cmd=%u (expected %u)" % (cmd, cmdin))


def get_master_version(conn):
	data, length = conn.command(CLTOMA_INFO, MATOCL_INFO)
	if length == 52:
		return (1, 4, 0)
	if length == 60:
		return (1, 5, 0)
	if length < 4:
		raise MFSCommunicationError("MATOCL_INFO too short: %u bytes" % length)
	return struct.unpack(">HBB", data[:4])


def supports_licence_mode(version):
	xyz, _ = version_convert(version)
	return (xyz >= (3, 0, 116) and xyz[0] == 3) or xyz >= (4, 27, 0)


def supports_licence_lstring(version):
	xyz, _ = version_convert(version)
	return xyz >= (4, 43, 0)


def get_licence(conn, version):
	# Returns a dict describing the licence, or None if the master refuses
	# (typical for Community Edition).
	if supports_licence_lstring(version):
		req = struct.pack(">B", 2)
		licver = 2
	elif supports_licence_mode(version):
		req = struct.pack(">B", 1)
		licver = 1
	else:
		req = None
		licver = 0

	try:
		data, length = conn.command(CLTOMA_LICENCE_DATA, MATOCL_LICENCE_DATA, req)
	except MFSCommunicationError:
		return {"licver": LICVER_CE, "lictypestr": "Community Edition (GPL2)"}

	if licver == 0 and length >= 15:
		lictype, v1, v2, v3, licmaxtime, licleft = struct.unpack(">BHBBLL", data[:13])
		lictypestr = ("DEMO" if lictype == 0 else
		              "COMMERCIAL" if lictype == 1 else
		              "FREE" if lictype == 2 else "UNKNOWN")
		pos = 13
		err = 0
		licissuer, pos, err = get_string_from_packet(data, pos, err)
		licuser, pos, err = get_string_from_packet(data, pos, err)
		if err:
			return None
		csmaxsize, csmaxcnt, quota = 0xFFFFFFFFFFFFFFFF, 0xFFFF, 0
		if length == 11 + pos:
			csmaxsize, csmaxcnt, quota = struct.unpack(">QHB", data[pos:])
		elif length == 10 + pos:
			csmaxsize, csmaxcnt = struct.unpack(">QH", data[pos:])
		elif length != pos:
			return None
		return {
			"licver": 0,
			"lictype": lictype, "lictypestr": lictypestr,
			"max_version": (v1, v2, v3),
			"licmaxtime": licmaxtime, "licleft": licleft,
			"licissuer": licissuer, "licuser": licuser,
			"csmaxsize": csmaxsize, "csmaxcnt": csmaxcnt, "quota": quota,
		}

	if licver == 1 and length >= 35:
		lictype, v1, v2, v3, licmaxtime, licleft, licmaxsize, currentsize = \
			struct.unpack(">BHBBLLQQ", data[:29])
		pos = 29
		err = 0
		lictypestr, pos, err = get_string_from_packet(data, pos, err)
		licissuer, pos, err = get_string_from_packet(data, pos, err)
		licuser, pos, err = get_string_from_packet(data, pos, err)
		licjson, pos, err = get_string_from_packet(data, pos, err)
		licid, pos, err = get_string_from_packet(data, pos, err)
		addinfo, pos, err = get_string_from_packet(data, pos, err)
		if err:
			return None
		return {
			"licver": 1, "lictype": lictype, "lictypestr": lictypestr,
			"max_version": (v1, v2, v3),
			"licmaxtime": licmaxtime, "licleft": licleft,
			"licmaxsize": licmaxsize, "currentsize": currentsize,
			"licid": licid, "licissuer": licissuer, "licuser": licuser,
			"licjson": licjson, "addinfo": addinfo,
		}

	if licver == 2 and length >= 47:
		lictype, v1, v2, v3, licmaxtime, licleft, licmaxsize, currentsize, licidnum = \
			struct.unpack(">BHBBLLQQQ", data[:37])
		pos = 37
		err = 0
		lictypestr, pos, err = get_longstring_from_packet(data, pos, err)
		licissuer, pos, err = get_longstring_from_packet(data, pos, err)
		licuser, pos, err = get_longstring_from_packet(data, pos, err)
		licjson, pos, err = get_longstring_from_packet(data, pos, err)
		addinfo, pos, err = get_longstring_from_packet(data, pos, err)
		if err:
			return None
		return {
			"licver": 2, "lictype": lictype, "lictypestr": lictypestr,
			"max_version": (v1, v2, v3),
			"licmaxtime": licmaxtime, "licleft": licleft,
			"licmaxsize": licmaxsize, "currentsize": currentsize,
			"licid": "%016X" % licidnum,
			"licissuer": licissuer, "licuser": licuser,
			"licjson": licjson, "addinfo": addinfo,
		}

	return None


def max_version_str(lic):
	if lic["licver"] == LICVER_CE:
		return "unlimited"
	v1, v2, v3 = lic["max_version"]
	if v1 == 0xFFFF:
		return "unlimited"
	if v2 == 255:
		return "%u.x" % v1
	if v3 == 255:
		return "%u.%u.x" % (v1, v2)
	return "%u.%u.%u" % (v1, v2, v3)


def expires_str(lic):
	if lic["licver"] == LICVER_CE:
		return "never"
	licleft = lic.get("licleft")
	licmaxtime = lic.get("licmaxtime")
	if licleft == 0xFFFFFFFF or licmaxtime == 0xFFFFFFFF:
		return "never"
	return time.asctime(time.localtime(licmaxtime))


def build_licence_rows(master_version, lic):
	rows = []
	if master_version is not None:
		strver, _ = version_str(master_version)
		rows.append(("Master version", strver))
	if lic is None:
		rows.append(("Licence", "unrecognized answer from master"))
	elif lic["licver"] == LICVER_CE:
		rows.append(("Edition", lic["lictypestr"]))
	else:
		if lic["licver"] >= 1 and lic.get("licid"):
			rows.append(("Licence ID", lic["licid"]))
		rows.append(("Issuer", lic.get("licissuer", "")))
		label = "User" if lic["licver"] == 0 else "Client"
		rows.append((label, lic.get("licuser", "")))
		extrajson = lic.get("licjson") or ""
		if extrajson:
			try:
				extra = json.loads(extrajson)
			except ValueError:
				extra = {}
			for k, v in extra.items():
				rows.append((str(k).lower(), str(v)))
		rows.append(("Type", lic["lictypestr"]))
		rows.append(("Max version", max_version_str(lic)))
		rows.append(("Expires", expires_str(lic)))
		if lic["licver"] == 0:
			if lic["csmaxsize"] == 0xFFFFFFFFFFFFFFFF:
				rows.append(("Max chunkserver size", "unlimited"))
			else:
				rows.append(("Max chunkserver size", humanize_number(lic["csmaxsize"], " ")))
			if lic["csmaxcnt"] == 0xFFFF:
				rows.append(("Max chunkservers", "unlimited"))
			else:
				rows.append(("Max chunkservers", str(lic["csmaxcnt"])))
		else:
			licmaxsize = lic["licmaxsize"]
			currentsize = lic["currentsize"]
			if (licmaxsize == 0xFFFFFFFFFFFFFFFF or (licmaxsize == 0 and currentsize == 0)):
				rows.append(("Max raw data usage", "unlimited"))
			else:
				rows.append(("Max raw data usage", humanize_number(licmaxsize, " ")))
			if lic.get("addinfo"):
				rows.append(("Additional info", lic["addinfo"]))
	return rows


def wrap_value(text, width):
	if len(text) <= width:
		return [text]
	out = []
	remaining = text
	while len(remaining) > width:
		cut = remaining.rfind(' ', 0, width + 1)
		if cut <= 0:
			cut = width
		out.append(remaining[:cut].rstrip())
		remaining = remaining[cut:].lstrip()
	if remaining:
		out.append(remaining)
	return out


def print_block(masters, rows, value_max_width=72):
	# Expand rows by soft-wrapping any value longer than value_max_width.
	# Continuation rows carry None as their key, rendered with a blank
	# key column and no colon to make wrapping visually obvious.
	expanded = []
	for k, v in rows:
		lines = wrap_value(v, value_max_width)
		expanded.append((k, lines[0]))
		for cont in lines[1:]:
			expanded.append((None, cont))

	if len(masters) == 1:
		header = "Master: %s" % masters[0]
	else:
		header = "Masters (%d): %s" % (len(masters), ", ".join(masters))
	keyw = max(len(k) for k, _ in expanded if k is not None)
	valw = max(len(v) for _, v in expanded)
	width = keyw + 3 + valw
	bar = "+" + "-" * (width + 2) + "+"
	print(header)
	print(bar)
	for k, v in expanded:
		if k is None:
			print("| %-*s   %-*s |" % (keyw, "", valw, v))
		else:
			print("| %-*s : %-*s |" % (keyw, k, valw, v))
	print(bar)


def parse_host_port(spec, default_port):
	# Accepts "host", "host:port", or bare IPv4 "1.2.3.4:port".
	if ":" in spec:
		h, _, p = spec.rpartition(":")
		try:
			return h, int(p)
		except ValueError:
			print("Invalid port in '%s'" % spec)
			sys.exit(2)
	return spec, default_port


def usage(prog):
	print("usage: %s [-h] [-v] [[host[:port] ...]" % prog)
	print("")
	print("Display MooseFS licence info from one or more master servers.")
	print("")
	print("  -h           print this help and exit")
	print("  -v           print version and exit")
	print("")
	print("Each positional argument is a master to query. The optional :port")
	print("suffix overrides the default master client port (%u)." % DEFAULT_MASTER_CLIENT_PORT)
	print("Failures on one host do not affect queries to the others.")


def main(argv):
	prog = argv[0]
	try:
		opts, args = getopt.gnu_getopt(argv[1:], "hv")
	except getopt.GetoptError as e:
		print("error: %s" % e)
		usage(prog)
		return 2
	for opt, _ in opts:
		if opt == "-h":
			usage(prog)
			return 0
		if opt == "-v":
			print("%s %s" % (prog,VERSION))
			return 0

	if not args:
		args = [DEFAULT_MASTERNAME]
#		print("error: no master hosts given")
#		usage(prog)
#		return 2

	# Expand each spec into one or more (host,port) targets via DNS.
	targets = []
	for spec in args:
		host, port = parse_host_port(spec, DEFAULT_MASTER_CLIENT_PORT)
		try:
			infos = socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM, socket.SOL_TCP)
		except socket.gaierror as e:
			targets.append((spec, host, port, None, "DNS error: %s" % e))
			continue
		ips = []
		for fam, stype, proto, _, sa in infos:
			if fam == socket.AF_INET and stype == socket.SOCK_STREAM:
				ips.append(sa[0])
		if not ips:
			targets.append((spec, host, port, None, "no IPv4 addresses resolved"))
			continue
		for ip in ips:
			targets.append((spec, host, port, ip, None))

	# Query every target, collecting results without printing.
	results = []  # list of (label, rows, error)
	for spec, host, port, ip, err in targets:
		display_host = host if ip is None or host == ip else "%s (%s)" % (host, ip)
		label = "%s:%u" % (display_host, port)
		if err is not None:
			results.append((label, None, err))
			continue
		try:
			conn = MFSConn(ip, port)
		except (socket.error, OSError) as e:
			results.append((label, None, "connection error: %s" % e))
			continue
		try:
			version = get_master_version(conn)
		except Exception as e:
			conn.close()
			results.append((label, None, "protocol error during version query: %s" % e))
			continue
		try:
			lic = get_licence(conn, version)
			results.append((label, build_licence_rows(version, lic), None))
		except Exception as e:
			results.append((label, None, "licence query failed: %s" % e))
		conn.close()

	# Group results: identical row sets collapse, identical errors collapse.
	licence_groups = []  # list of (rows, [labels])
	error_groups = []    # list of (err, [labels])
	for label, rows, err in results:
		if err is not None:
			for i, (e, labels) in enumerate(error_groups):
				if e == err:
					labels.append(label)
					break
			else:
				error_groups.append((err, [label]))
		else:
			key = tuple(rows)
			for existing, labels in licence_groups:
				if existing == key:
					labels.append(label)
					break
			else:
				licence_groups.append((key, [label]))

	total = len(results)
	first_block = True

	def sep():
		nonlocal first_block
		if not first_block:
			print("")
		first_block = False

	# Single-group, all-success, multi-master case: announce shared licence.
	if (len(licence_groups) == 1 and not error_groups and total > 1):
		print("All %d masters share the same licence." % total)
		print("")
		rows, labels = licence_groups[0]
		print_block(labels, rows)
	else:
		if len(licence_groups) > 1:
			covered = sum(len(labels) for _, labels in licence_groups)
			print("Found %d distinct licences across %d masters." % (len(licence_groups), covered))
			print("")
		for rows, labels in licence_groups:
			sep()
			print_block(labels, rows)
		for err, labels in error_groups:
			sep()
			print_block(labels, [("Error", err)])

	return 1 if error_groups else 0


if __name__ == "__main__":
	sys.exit(main(sys.argv))
