# 升级wordpress新版本之后运行一次修复发送邮件链接不正确的问题

# 作者：雪山凌狐
# 日期：2020-01-03
# 版本：1.1
# 官网地址：http://www.xueshanlinghu.com

# 平台：python3

import sys
import os

#==============================配置============================
WP_LOGIN_LINE = "$message .= '<' . network_site_url( \"wp-login.php?action=rp&key=$key&login=\" . rawurlencode( $user_login ), 'login' ) . \">\\r\\n\";"
WP_LOGIN_OK = "$message .= network_site_url( \"wp-login.php?action=rp&key=$key&login=\" . rawurlencode($user_login), 'login') . \"\\r\\n\";"

PLUGGABLE_LINE = "$message .= '<' . network_site_url( \"wp-login.php?action=rp&key=$key&login=\" . rawurlencode($user->user_login), 'login') . \">\\r\\n\\r\\n\";"
PLUGGABLE_OK = "$message .= network_site_url( \"wp-login.php?action=rp&key=$key&login=\" . rawurlencode($user->user_login), 'login') . \"\\r\\n\\r\\n\";"
#==============================================================

# 目录_处理，将右边加上\
def deal_folder(folder):
	if folder[-1] != "\\":
		new_folder = folder + "\\"
	else:
		new_folder = folder
	
	return new_folder

# 获取文件夹的所有文件
def get_all_file(folder):
	if not os.path.exists(folder):
		return []
	else:
		re = []
		r = os.listdir(folder)
		for i in r:
			if os.path.isfile(folder + i):
				re.append(i)
		return re

# 获取wp-login文件路径		
def get_wp_login(folder):
	filelist = get_all_file(folder)
	for i in filelist:
		if i == "wp-login.php":
			return folder + "wp-login.php"
	return None

# 修复wp-login文件
def edit_wp_login(path):
	re = False
	f = open(path, 'r', encoding = 'utf-8')
	lines = f.readlines()
	f.close()
	f = open(path, 'w', encoding = 'utf-8')
	for line in lines:
		if not re:
			if line.find(WP_LOGIN_LINE) > -1:
				re = True
				line = line.replace(WP_LOGIN_LINE,WP_LOGIN_OK)
		f.write(line)	
	f.close()
	return re

# 获取pluggable文件路径
def get_pluggable(folder):
	if not os.path.exists(folder + "wp-includes\\"):
		return None
	filelist = get_all_file(folder + "wp-includes\\")
	for i in filelist:
		if i == "pluggable.php":
			return folder + "wp-includes\\pluggable.php"
	return None

# 修复pluggable文件
def edit_pluggable(path):
	re = False
	f = open(path, 'r', encoding = 'utf-8')
	lines = f.readlines()
	f.close()
	f = open(path, 'w', encoding = 'utf-8')
	for line in lines:
		if not re:
			if line.find(PLUGGABLE_LINE) > -1:
				re = True
				line = line.replace(PLUGGABLE_LINE,PLUGGABLE_OK)
		f.write(line)	
	f.close()
	return re

# 主入口函数	
if __name__ == '__main__':
	
	# 判定并获取网站根目录路径参数
	if len(sys.argv) < 2:
		print("请输入网站根目录路径参数")
		exit()
	else:
		print("您设定的网站根目录为：")
		folder = deal_folder(sys.argv[1])
		print(folder)
		print("=" * 40)
	
	# wp-login.php文件修改
	print("正在寻找wp-login.php文件中...")
	wp_login_path = get_wp_login(folder)
	if wp_login_path:
		print("文件已找到，修改中...")
		re = edit_wp_login(wp_login_path)
		if re:
			print("修改成功！")
		else:
			print("找不到该语句，请检查！")
	else:
		print("wp-login文件不存在！！！请检查！！！")
	
	print("=" * 40)
	
	# pluggable.php文件修改	
	print("正在寻找pluggable.php文件中...")
	pluggable_path = get_pluggable(folder)
	if pluggable_path:
		print("文件已找到，修改中...")
		re = edit_pluggable(pluggable_path)
		if re:
			print("修改成功！")
		else:
			print("找不到该语句，请检查！")
	else:
		print("pluggable文件不存在！！！请检查！！！")
	
	print("=" * 40)
	print("运行完毕，感谢您的使用！")
	# 用来暂停进程，查看结果后再关闭
	input("Please press Enter button to exit...")