from DrissionPage import ChromiumPage, ChromiumOptions
from requests.cookies import RequestsCookieJar
from retry import retry
import time
import os
import pickle





page = ChromiumPage()
@retry(tries=50, delay=1)
def get_cookie():
    tab = page.new_tab(url='https://www.douyin.com/user/self')

    time.sleep(1)

    while True:
        if tab.ele('x://div[@data-e2e="user-detail"]//h1').text != '未登录':
            print('登录成功')
            break
        if tab.ele('x://*[text()="编辑资料"]').text != '':
            print('登录成功')
            break
        time.sleep(1)

    title = tab.title
    turl = tab.ele('x://link[@rel="canonical"]').attr('href')
    # print(title)
    print(turl)
    filename = title.replace('的抖音 - 抖音', '')
    if filename == "":
        filename = tab.ele('x://div[@data-e2e="user-detail"]//h1').text
    cookies = tab.cookies()
    cookie_jar = RequestsCookieJar()
    for cookie in cookies:
        cookie_jar.set(cookie['name'], cookie['value'], domain=cookie['domain'])
    turl = turl.replace('https://www.douyin.com/user/', '')
    cntext = f'{turl}|{filename}\r'
    if os.path.isfile('cookie_name.txt'):
        with open('cookie_name.txt', 'rb') as fp:
            key = fp.read().decode('utf-8')
        # 分割字符串得到列表
        klist = key.split('\r')

        seen = set()
        unique_klist = [x for x in klist if not (x in seen or seen.add(x))]
        cntext += '\r'.join(unique_klist)

    with open('./cookie_name.txt', 'w', encoding='utf-8') as f:
        f.write(cntext)
    with open(f'./{turl}.pkl', 'wb') as f:
        pickle.dump(cookie_jar, f)
    print('成功')
    input(111111)
    
             



    



if __name__ == '__main__':
    get_cookie()
    page.quit()