from flask import Flask, request, jsonify
import os
import requests

app = Flask(__name__)

# 定义全局字符串变量
str_values = ''
str_uvalues = ''
# 从环境变量中获取密码
password = os.environ.get('PASSWORD', '123456')

# 定义一个函数，接受 Cookie-Values 和要保留的键的列表作为参数
def filter_cookie_values(cookie_values, keep_keys):
    new_cookie_values = ""
    pairs = [pair.strip() for pair in cookie_values.split(';')]
    for pair in pairs:
        if '=' in pair:
            key, value = pair.split('=', 1)
            key = key.strip()
            value = value.strip()
            if key in keep_keys:
                new_cookie_values += f'{key}={value}; '
    new_cookie_values = new_cookie_values[:-2]
    return new_cookie_values

@app.route('/SET', methods=['POST'])
def set_value():
    method = request.method
    pwd = request.args.get('pwd')
    if not pwd or pwd!= password:
        return 'Invalid password', 401
    keep_keys = ["_U", 
                "MUID",
                'KievRPSSecAuth',
                'cct',
                'buid',
                'ak_bmsc',
                'bm_sv',
                'cl_dtrs',
               '_RwBf',
               'SRCHHPGUSR',
               'WLS']
    keep_keys_u = ["_U",
                "WLS"]
    cookie_values = request.headers.get('Cookie-Values')
    set_value = filter_cookie_values(cookie_values, keep_keys)
    get_u_value = filter_cookie_values(cookie_values, keep_keys_u)
    if set_value:
        str_values = set_value
        if get_u_value and get_u_value not in str_uvalues:
            str_uvalues += ';' + get_u_value
        return 'Set value successfully'
    else:
        return 'No Cookie-Values in header', 400

@app.route('/GET', methods=['GET'])
def get_value():
    method = request.method
    pwd = request.args.get('pwd')
    if not pwd or pwd!= password:
        return 'Invalid password', 401
    result = {'result': {'cookies': str_values}}
    return jsonify(result)

@app.route('/CLS', methods=['GET', 'POST'])
def clear_value():
    method = request.method
    pwd = request.args.get('pwd')
    if not pwd or pwd!= password:
        return 'Invalid password', 401
    replaced_str = str_uvalues.replace(';', '<br>')
    str_values = ''
    str_uvalues = ''
    return 'Clear value successfully' + '\n' + replaced_str

@app.route('/HisU', methods=['GET', 'POST'])
def history_u():
    method = request.method
    pwd = request.args.get('pwd')
    if not pwd or pwd!= password:
        return 'Invalid password', 401
    replaced_str = str_uvalues.replace(';', '<br>')
    return 'Ukey History:' + '\n' + replaced_str

@app.route('/', methods=['GET', 'POST'])
def root():
    return 'Please visit /SET /GET or /CLS with?pwd=xxxxxx'

@app.route('/q', methods=['GET'])
def q_route():
    url = request.args.get('url')
    if not url:
        return jsonify({'error': 'URL is required'}), 400
    keep_keys = ["_U", 
                "MUID",
                'KievRPSSecAuth',
                'cct',
                'buid',
                'ak_bmsc',
                'bm_sv',
               '_RwBf',
               'SRCHHPGUSR',
               'WLS']
    response = requests.get(url, headers={
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0',
        'Accept-Language': 'zh-CN,zh;q=0.9'
    })
    set_cookie_array = response.headers.get('set-cookie', [])
    if set_cookie_array:
        cookies = [cookie.split(';')[0] for cookie in set_cookie_array]
        cookies_str = '; '.join(cookies)
    else:
        cookies_str = ''
    set_value = filter_cookie_values(cookies_str, keep_keys)
    result = {'result': {'cookies': set_value}}
    return jsonify(result)

if __name__ == '__main__':
    app.run(port=7860)