#!/usr/bin/python3
# coding=utf-8
##################################################################
#              ____     _     _ __  __                 
#             / __/__  (_)___(_) /_/ /  ___  ___  ___ _
#            _\ \/ _ \/ / __/ / __/ /__/ _ \/ _ \/ _ `/
#           /___/ .__/_/_/ /_/\__/____/\___/_//_/\_, / 
#              /_/                              /___/  
# 
# Copyright (c) 2026 Chongqing Spiritlong Technology Co., Ltd.
# All rights reserved.
# @author	arthuryang
# @brief	服务器相关工具
##################################################################

import	os
import	re
import	time
import	datetime
import	string
import	secrets
import	random
import	SpiritLong_utility
import	SpiritLong_email

## 初始化服务器，不包含修改密钥/密码的操作
#	hostanme	主机HOSTNAME：大写字母、数字和横线
#	address		访问该主机的IP地址或者域名
#	port_SSH	SSH端口，None将生成随机端口
#	port_MySQL	MySQL端口
def server_initialize(server_type, hostname, address, password_root, MySQL_password_spiritlong=None, port_SSH=None, port_MySQL=None):
	# 本地脚本
	local_file	= os.path.join(os.path.dirname(__file__), f"ECS_initialize_{server_type}.sh")
	print(local_file)
	if not os.path.exists(local_file):
		print("没有此服务器类型的初始化脚本")
		return
	
	if port_SSH is None:
		# 默认SSH使用50000-60000范围内的随机端口
		port_SSH	= random.randint(50000, 60000)
	
	if port_MySQL is None:
		# 默认MySQL使用3500-3999范围内的随机端口
		port_MySQL	= random.randint(3500, 3999)
	
	if MySQL_password_spiritlong is None:
		MySQL_password_spiritlong	= SpiritLong_utility.generate_random_password(alphanumeric=True)

	print(f"初始化主机	{hostname}")
	print(f"地址		{address}")
	print(f"SSH端口		{port_SSH}")
	print(f"MySQL端口	{port_MySQL}")

	# 执行脚本进行服务器初始化
	remote_file	= '/tmp/ECS_initialize.sh'

	# 新初始化后会是此默认SSH端口
	port		= 22

	print(f"--------------- 主机{hostname}初始化 ---------------")

	# 尝试连接到服务器
	try_count	= 10
	print(f"尝试连接服务器{address}，root密码为{password_root}")
	while(True):
		result	= SpiritLong_utility.execute_on_server(address, port, 'root', password_root, 'echo OK')
		if result=='OK\n' :
			print('\t服务器root用户登录成功')
			break
		else:
			print("\t等待10秒再进行下一次尝试")
			time.sleep(10)
			try_count	-= 1
		if try_count==0:
			break
	if try_count==0:
		print("多次重试后仍无法连接")
		return

	# 上传脚本进行初始化
	print(f"服务器开始初始化，预计要等待4分钟左右，当前时刻{datetime.datetime.now()}")
	SpiritLong_utility.upload_on_server(address, port, 'root', password_root, local_file, remote_file)
	echo_string	= "server initilaized done."
	log_file	= "/root/server_initialize.log"
	result	= SpiritLong_utility.execute_on_server(address, port, 'root', password_root, f'''
		HOSTNAME={hostname} PORT_SSH={port_SSH} PORT_MYSQL={port_MySQL} MYSQL_SPIRITLONG_PASSWORD={MySQL_password_spiritlong} . {remote_file} > {log_file} 2>&1 && rm {remote_file} && echo "{echo_string}"
''')
	print("服务器初始化"+("成功" if result==f'{echo_string}\n' else "失败"))

	return {
		'hostname'			: hostname,
		'address'			: address,
		'port'				: port_SSH,
		'user'				: 'root',
		'key'				: None,
		'port_MySQL'			: port_MySQL,
		'MySQL_password_spiritlong'	: MySQL_password_spiritlong,
		'password_root'			: password_root,
	}
	
## 在本地增加一个session，设置主机root密钥/root密码
#	username		在root公钥备注中要用到
#	session			{hostname, address, port, user, key_file, {key_type, public, private, comment}}
#	password_root		现用的root密码
#	update_password_root	更新root密码为随机密码
def server_root_authorization_setup(username, session, password_root=None, update_password_root=False, password_length=20):
	# 增加session
	session_name	= SpiritLong_utility.utility_SSH._get_session_name_from_session(session)
	sessions	= {session_name:session}
	SpiritLong_utility.add_sessions(sessions)

	print(f"--------------- 主机{session['hostname']}设置root权限 ---------------")

	# 远程执行
	SpiritLong_utility.execute_on_server(session['address'], session['port'], 'root', password_root, f'''
		mkdir -p /root/.ssh
		echo "{session['key']['key_type']} {session['key']['public']} {username}" >> /root/.ssh/authorized_keys
''')

	# 测试连接
	SpiritLong_utility.test_sessions_connectivity(sessions)
	
	# 更新密码，注意bash中单引号将把$视为字符，而双引号将会视为shell变量
	if update_password_root:
		password	= SpiritLong_utility.update_password_on_session(session_name, 'root', password_length=password_length)
		print(f"新的密码\t{password}")
	else:
		password	= password_root
	
	session['password']	= password

	return session

## 执行数据库操作
def mysql_execute_on_session(session_name, command, print_result=False):
	# 命令预处理：反斜杠、反引号和美元号需要转义，否则bash命令中的双引号会错误理解
	command	= re.sub(r'\\', r'\\', command)
	command	= re.sub(r'`', r'\\`', command)
	command	= re.sub(r'\$', r'\\$', command)
	result	= SpiritLong_utility.execute_on_session(session_name, f'''mysql -e "{command}"  --default-character-set=utf8mb4''')
	if result is None:
		return
	result	= result.strip()
	if result and print_result:
		print(result)
	return result

## 执行SELECT，返回字典列表。只支持单表查询
#	session_name		session名称
#	database		数据库名称
#	table			表
#	fields			字段列表，可以传入单个字段名称的字符串，不要使用AS，不可使用*
#	where			WHERE语句的后半截
#	order_by		ORDER BY语句的后半截
#	limit_n			LIMIT语句的数量
#	print_SQL_string	打印所用的SQL语句，以便调试
# 返回结果的字典列表[{field:value}]，没有结果将返回[]
def mysql_select_on_session(session_name, database, table, fields, where=None, order_by=None, limit_n=None, print_SQL_string=False):
	select_result	= []

	# 字段列表，每个字段都要用反引号包围
	if not fields:
		print("必须指定fileds")
		return select_result
	if isinstance(fields, str):
		if fields.strip()=='*':
			print("不支持使用*")
			return select_result
		else:
			fields_string	= f"QUOTE(`{fields}`)"
			fields		= [fields]
	else:
		fields_string	= f" {','.join([f"QUOTE(`{f}`)" for f in fields])} "


	# WHERE
	where_string	= f" WHERE {where} " if where else ''

	# ORDER
	order_string	= f" ORDER BY `{order_by}` " if order_by else ''

	# LIMIT
	limit_string	= f" LIMIT {limit_n} " if limit_n else ''

	SQL_string	= f"SELECT {fields_string} FROM `{database}`.`{table}` {where_string} {order_string} {limit_string}"
	result		= mysql_execute_on_session(session_name, SQL_string)

	if result:
		# 转换为字典列表
		records	= result.split('\n')
		# 第一行是字段名
		keys	= re.findall(r'QUOTE\(`(.*?)`\)', records[0])
		#  比较keys和fields
		common		= [v for v in keys if v in fields]
		different	= [v for v in keys+fields if v not in common]
		if different and fields!='*':
			# 有差异元素
			print(f"返回的字段列表和给定的字段列表有差异：{different}")
			print(result)
			return select_result
		for r in records[1:]:
			# 需要手动分解：考虑NULL是没有引号的，单引号将会以转义方式出现\\
			# 去掉头尾可能的单引号
			if r.startswith("'"):
				r	= r[1:]
			if r.endswith("'"):
				r	= r[:-1]
			# 只有NULL和引号包含的两种情况，有4种分隔条件
			values	= re.split(r"^(?<=NULL)|(?=NULL)$|'\t'|'\t(?=NULL)|(?<=NULL)\t'|(?=NULL)\t(?<=NULL)", r)
			select_result.append({fields[i]:values[i] for i in range(len(fields))})
	return select_result

## 创建/更新数据库，以及数据库专用用户
#	session_name	session名称
#	database_name	数据库名
#	schema_file	数据库结构定义文件，默认为None
#	data_file	数据库数据，默认为None
def add_mysql_database(session_name, database_name, schema_file=None, data_file=None):

	print(f"创建数据库{database_name}")
	password		= SpiritLong_utility.generate_random_password(alphanumeric=True)
	password_localhost	= SpiritLong_utility.generate_random_password(alphanumeric=True)
	
	# 创建数据库
	mysql_execute_on_session(session_name, f"CREATE DATABASE IF NOT EXISTS `{database_name}`;")
	
	# 检查用户{database_name@%是否存在
	result	= mysql_select_on_session(session_name, database_name, 'USER', 'ID', where=f"user={database_name} AND host='%'")
	if result:
		# 删除此用户{database_name@%
		SQL_string	= f"DROP USER {database_name};"
		result		= mysql_execute_on_session(session_name, SQL_string)
		
	# 添加用户并授权此数据库全部权限
	SQL_string	= f'''CREATE USER '{database_name}'@'%' IDENTIFIED BY '{password}';
		GRANT ALL PRIVILEGES ON `{database_name}`.* TO {database_name} WITH GRANT OPTION; 
		FLUSH PRIVILEGES;
	'''
	result		= mysql_execute_on_session(session_name, SQL_string)
	print(f"添加数据库用户{database_name}，密码\t{password}")

	# 导入数据库结构
	if schema_file and os.path.exists(schema_file):
		print(f"从{schema_file}导入数据库结构")
		# 数据库表
		remote_filename	= "/tmp/schema.xlsx"
		# 上传结构表文件
		SpiritLong_utility.upload_on_session(session_name, schema_file, remote_filename)
		# 执行导入脚本
		result	= SpiritLong_utility.execute_on_session(session_name, f'''
			PY_SCRIPT=/tmp/db.py
			cat > $PY_SCRIPT <<EOF
#!/usr/bin/python3
import SpiritLong_database

db      = SpiritLong_database.Database({{
        "host"          : "localhost",
        "user"          : "root",
        "passwd"        : "",
        "database"      : "{database_name}",
        "charset"       : "utf8mb4",
}})

db.import_excel_schema("{remote_filename}")
EOF
			chmod +x $PY_SCRIPT
			$PY_SCRIPT
			rm $PY_SCRIPT
			rm {remote_filename}
''')
		print(result)

	# 导入数据
	if data_file and os.path.exists(data_file):
		print(f"从文件导入数据库数据\t{data_file}")
		# 数据库表
		remote_filename	= "/tmp/data.xlsx"
		# 上传数据文件
		SpiritLong_utility.upload_on_session(session_name, data_file, remote_filename)
		# 执行导入脚本
		result	= SpiritLong_utility.execute_on_session(session_name, f'''
			PY_SCRIPT=/tmp/db.py
			cat > $PY_SCRIPT <<EOF
#!/usr/bin/python3
import SpiritLong_database

db      = SpiritLong_database.Database({{
        "host"          : "localhost",
        "user"          : "root",
        "passwd"        : "",
        "database"      : "{database_name}",
        "charset"       : "utf8mb4",
}})

db.import_excel_data("{remote_filename}")
EOF
			chmod +x $PY_SCRIPT
			$PY_SCRIPT
			rm $PY_SCRIPT
			rm {remote_filename}
''')
		print(result)

## 配置此系统使用mysql管理用户
#	session_name		连接到主机的sessnion名称
#	pause_before_pam	在设置PAM之前可以停一下，立即开一个终端连接到该主机，防止PAM配置失败，还有手动操作的机会
def setup_linux_user_via_mysql(session_name, pause_before_pam=False):
	sessions	= SpiritLong_utility.get_sessions()
	result	= [s for n, s in sessions.items() if n==session_name and s['user']=='root']
	if not result:
		return 
	
	session	= result[0]

	print(f"--------------- 主机{session['hostname']}设置为通过mysql数据库管理用户 ---------------")


	# 创建以HOSTNAME命名的数据库来管理用户
	database_name	= session['hostname']
	schema_file	= os.path.join(os.path.dirname(__file__), 'database_schema.xlsx')
	data_file	= os.path.join(os.path.dirname(__file__), 'database_data.xlsx')
	add_mysql_database(session_name, database_name, schema_file, data_file)
	
	# 添加数据库用户PAM和PAM_root
	username		= "'PAM'@'localhost'"
	password_PAM		= SpiritLong_utility.generate_random_password(alphanumeric=True)
	SQL_string	= f'''	CREATE USER {username} IDENTIFIED BY '{password_PAM}'; 
				GRANT SELECT(ID,USER,NAME,GROUP_ID,DELETE_TIMESTAMP) ON `{database_name}`.`USER` TO {username};
				GRANT SELECT ON `{database_name}`.`GROUP` TO {username};
				GRANT SELECT ON `{database_name}`.`_USER_GROUP` TO {username};
				FLUSH PRIVILEGES;
			'''
	result	= mysql_execute_on_session(session_name, SQL_string)

	username		= "'PAM_root'@'localhost'"
	password_PAM_root	= SpiritLong_utility.generate_random_password(alphanumeric=True)
	result	+= SpiritLong_utility.execute_on_session(session_name, f'''
		mysql -e "CREATE USER {username} IDENTIFIED BY '{password_PAM_root}'; 
			GRANT ALL PRIVILEGES ON \\`{database_name}\\`.* TO {username} WITH GRANT OPTION;
			FLUSH PRIVILEGES;"  --default-character-set=utf8mb4
	''')
	if result.strip():
		print(result)

	print("新增数据库用户：")
	print(f"PAM		{password_PAM}")
	print(f"PAM_root	{password_PAM_root}")

	# 安装pam-mysql和libnss-mysql(Debian13提供)
	result	= SpiritLong_utility.execute_on_session(session_name, f'''apt-get -y install libpam-mysql libnss-mysql finger''')
	if result!='\n':
		print(result)
	
	# 配置pam-mysql
	result	= SpiritLong_utility.execute_on_session(session_name, f'''
		cat > /etc/pam-mysql.conf <<EOF
users.host              = localhost
users.database          = {database_name}
users.db_user           = PAM_root
users.db_passwd         = {password_PAM_root}
users.table             = USER
users.user_column       = USER
users.password_column   = PASSWORD
users.password_crypt    = 1
users.use_md5          = yes
users.where_clause     = DELETE_TIMESTAMP IS NULL
EOF
	''')

	if pause_before_pam:
		# 可以在此暂停
		answer	= input("接下来要配置权限，可能会导致无法登录！务必立即开启至少一个root终端以防万一。继续？输入N/n将中止").strip()
		if answer.startswith(('n', 'N')):
			return

	# 设置pam.d中公用部分的认证方式
	for item in ['account', 'password', 'auth', 'session', 'session-noninteractive']:
		start_string	= item if item!='session-noninteractive' else 'session'
		result	+= SpiritLong_utility.execute_on_session(session_name, f'''
			sed -i '0,/^$/s&^$&{start_string} sufficient pam_mysql.so config_file=/etc/pam-mysql.conf&' /etc/pam.d/common-{item}
		''')	  
	if result.strip():
		print(result)

	# 设置libnss-mysql：配置文件和设置nsswitch认证方式，注意`和$需要用\转义，字符串中则需要用\\来表示\`
	result	= SpiritLong_utility.execute_on_session(session_name, f'''
		cat > /etc/libnss-mysql.cfg <<EOF
getpwnam	SELECT \\`USER\\`,'x',\\`ID\\`,\\`GROUP_ID\\`,\\`NAME\\`,concat('/home/',\\`USER\\`),'/bin/bash'	FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`='%1\\$s' LIMIT 1
getpwuid	SELECT \\`USER\\`,'x',\\`ID\\`,\\`GROUP_ID\\`,\\`NAME\\`,concat('/home/',\\`USER\\`),'/bin/bash'	FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`ID\\`='%1\\$u' LIMIT 1
getspnam	SELECT \\`USER\\`,'x','0','0','99999','0','0','-1','0'				FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`='%1\\$s' LIMIT 1
getpwent	SELECT \\`USER\\`,'x',\\`ID\\`,\\`GROUP_ID\\`,\\`NAME\\`,concat('/home/',\\`USER\\`),'/bin/bash'	FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\` is NOT NULL
getspent	SELECT \\`USER\\`,'x','0','0','99999','0','0','-1','0'				FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\` is NOT NULL
getgrnam	SELECT \\`GROUP\\`,'x',\\`ID\\`								FROM \\`GROUP\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`GROUP\\`='%1\\$s' LIMIT 1
getgrgid	SELECT \\`GROUP\\`,'x',\\`ID\\`								FROM \\`GROUP\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`ID\\`='%1\\$u' LIMIT 1
getgrent	SELECT \\`GROUP\\`,'x',\\`ID\\`								FROM \\`GROUP\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\` is NOT NULL
memsbygid	SELECT \\`USER_ID\\`	FROM \\`_USER_GROUP\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`GROUP_ID\\`='%1\\$u' UNION SELECT \\`ID\\`	FROM \\`USER\\`	WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`.\\`GROUP_ID\\`='%1\\$u'
gidsbymem	SELECT \\`_USER_GROUP\\`.\\`GROUP_ID\\` FROM \\`_USER_GROUP\\` JOIN \\`USER\\` ON \\`USER\\`.\\`ID\\`=\\`_USER_GROUP\\`.\\`USER_ID\\` WHERE \\`_USER_GROUP\\`.\\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`.\\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`.\\`USER\\`='%1\\$s' UNION SELECT \\`GROUP_ID\\` FROM \\`USER\\` WHERE \\`DELETE_TIMESTAMP\\` is NULL AND \\`USER\\`='%1\\$s'

host		localhost
database	{database_name}
username	PAM
password	{password_PAM}
EOF
		cat > /etc/libnss-mysql-root.cfg <<EOF
username        PAM_root
password        {password_PAM_root}
EOF
		sed -i 's/systemd/mysql systemd/' /etc/nsswitch.conf
	''')

## 添加系统用户
#	session_name
#	user		{ID, NAME, USER, GROUP, IDENTITY, MOBILE, EMAIL, REMARK}, 其中NAME/USER/GROUP必填，ID,IDENTITY,MOBILE,EMAIL, REMARK为可选
def add_system_user(session_name, user):
	session	= SpiritLong_utility.get_session_from_name(session_name)
	hostname	= session['hostname']
	database_name	= hostname

	# 检查基本信息是否齐全
	key_pass	= True
	for key in ['NAME', 'USER', 'GROUP']:
		if key not in user:
			print(f"user字典中必须有{key}")
			key_pass	= False
	if not key_pass:
		return
	
	# 检查该组的ID
	result	= mysql_select_on_session(session_name, database_name, 'GROUP', 'ID', where=f"`GROUP`='{user['GROUP']}'")
	if not result:
		print(f"没有找到组{user['GROUP']}")
		return
	user['GROUP_ID']	= result[0]['ID']
	del user['GROUP']
	
	set_values	= []
	for key in user:
		set_values.append(f"`{key}`='{user[key]}'")
		
	# 检查该用户是否存在
	result	= mysql_select_on_session(session_name, database_name, 'USER', ['ID', 'DELETE_TIMESTAMP'], where=f"USER='{user['USER']}'")
	if result:
		if result[0]['DELETE_TIMESTAMP']:
			# 已经标记删除，恢复并更新信息
			print(f"用户{user['USER']}曾经存在，尝试恢复并更新信息")
			set_values.append(f"`DELETE_TIMESTAMP`=NULL")
			SQL_string	= f"UPDATE `{database_name}`.USER SET {','.join(set_values)} WHERE ID={result[0]['ID']};"
		else:
			print(f"用户{user['USER']}已经存在")
			return
	else:
		SQL_string	= f"INSERT INTO `{database_name}`.USER SET {','.join(set_values)};"
	
	# 增加/恢复该用户
	result		= mysql_execute_on_session(session_name, SQL_string)

## 删除/禁用系统用户
#	session_name
#	user		{ID, NAME, USER}，只要有其中一项即可，按ID/USER/NAME的优先级（仅NAME唯一时才可删除）
def delete_system_user(session_name, user):
	
	if 'ID' in user:
		where	= f"ID={user['ID']}"
	elif 'USER' in user:
		where	= f"USER='{user['USER']}'"
	elif 'NAME' in user:
		# 检查NAME是否唯一
		# 注意session_name必定和hostname是一样的
		result	= mysql_select_on_session(session_name, session_name, 'USER', 'ID', where=f"NAME='{user['NAME']}'")
		if len(result)==0:
			# 没找到
			print(f"没找到NAME为'{user['NAME']}'的用户")
			return
		elif len(result)==1:
			# 有唯一的
			where	= f"ID={result[0]['ID']}"
		else:
			print(f"有{len(result)}个NAME为'{user['NAME']}'的用户")
			return
	else:
		print("无法确定用户信息，需要ID/USER/NAME之一（仅NAME唯一时才可删除）")
	
	SQL_string	= f"UPDATE `{session_name}`.USER SET `DELETE_TIMESTAMP`=now() WHERE {where}"
	result		= mysql_execute_on_session(session_name, SQL_string)
	print(result)

## 在指定主机上初始化用户目录，并增加samba用户，选配邮件发送
#	session_name	
#	username	用户名
#	database_name	对于从数据库，需要使用主数据库的名称，None表示和session_name一样
#	mail_sender	{host, port, user, password}
def initialize_system_user_on_session(session_name, username, database_name=None, mail_sender=None):
	if not database_name:
		database_name	= session_name

	# 定位用户记录
	result	= mysql_select_on_session(session_name, database_name, 'USER', ['USER', 'EMAIL', 'NAME'], f"USER='{username}'")

	if not result:
		# 没查询到结果
		print(f"未找到用户{username}的记录")
		return
	user	= result[0]

	# 设置用户密码
	password	= SpiritLong_utility.generate_random_password(alphanumeric=True)
	SpiritLong_utility.update_password_on_session(session_name, username, password)

	# 设置本地home目录，设置samba密码（和用户密码相同
	result	= SpiritLong_utility.execute_on_session(session_name, f'''
		mkdir -p /home/{username}
		cp /etc/skel/.* /home/{username}/
		chown {username} /home/{username} -R
		echo "{password}\n{password}\n" | smbpasswd -a {username} -s
		SAMBA_PASSWORD_FILE=/home/{username}/user_password
		cat > $SAMBA_PASSWORD_FILE <<EOF
initial user/samba password:
{password}
EOF
		chown {username} $SAMBA_PASSWORD_FILE
		chmod 600 $SAMBA_PASSWORD_FILE
	''')

	# 发送邮件
	if mail_sender and user['EMAIL']:
		# 寻找用户名
		try:
			subject	= f"{database_name}系统账号生效通知" 
			content	= f'''{user['NAME'] if user['NAME'] else username}:

{database_name}系统账号{username}已经生效，用户密码和网上邻居都为

{password}

----------------------
{database_name}系统 致
{datetime.datetime.now():%Y-%M-%d}
'''
			email	= SpiritLong_email.Email(mail_sender['host'], mail_sender['port'], mail_sender['user'], mail_sender['password'])
			# 发送
			email.send_mail(user['EMAIL'], subject, content, [])
		except Exception as e:
			print(f"发送邮件失败：{e}")
	return 

if __name__=='__main__':
	pass