# -*- coding:utf-8 -*-
import os
import json
import base64
import sqlite3
from win32crypt import CryptUnprotectData
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

#  pip install pywin32
#  pip install cryptography
#  文檔:https://source.chromium.org/chromium/chromium/src/+/master:components/os_crypt/os_crypt_win.cc?q=OSCrypt&ss=chromium

class Chrome:
  def __init__(self):
    self.local_state = os.environ['LOCALAPPDATA'] + r'\Google\Chrome\User Data\Local State'
    self.cookie_path = os.environ['LOCALAPPDATA'] + r"\Google\Chrome\User Data\Default\Login Data"

  def get_key(self):
    with open(self.local_state, 'r', encoding='utf-8') as f:
      base64_encrypted_key = json.load(f)['os_crypt']['encrypted_key']
    encrypted_key_with_header = base64.b64decode(base64_encrypted_key)
    # 去掉開(kāi)頭的DPAPI
    encrypted_key = encrypted_key_with_header[5:]
    key_ = CryptUnprotectData(encrypted_key, None, None, None, 0)[1]
    return key_

  @staticmethod
  def decrypt_string(key, secret, salt=None):
    """
    解密
    """
    # 去掉'v10'
    nonce, cipher_bytes = secret[3:15], secret[15:]
    aes_gcm = AESGCM(key)
    return aes_gcm.decrypt(nonce, cipher_bytes, salt).decode('utf-8')

  @staticmethod
  def encrypt_string(key, data, salt=None):
    """
    加密
    """
    aes_gcm = AESGCM(key)
    prefix = "v10".encode("utf-8")
    # 隨機(jī)生成12位字符串,拼接"v10" 共15位
    nonce = os.urandom(12)
    cipher_bytes = data.encode("utf-8")
    return prefix + nonce + aes_gcm.encrypt(nonce, cipher_bytes, salt)

  def get_password(self, host):
    sql = f"select username_value,password_value from logins where signon_realm ='{host}';"
    with sqlite3.connect(self.cookie_path) as conn:
      cu = conn.cursor()
      res = cu.execute(sql).fetchall()
      cu.close()
      result = []
      key = self.get_key()

      for name, encrypted_value in res:

        if encrypted_value[0:3] == b'v10' or encrypted_value[0:3] == b'v11':
          password = self.decrypt_string(key, encrypted_value)
        else:
          password = CryptUnprotectData(encrypted_value)[1].decode()
        result.append({'user_name': name, 'password': password})
      return result

  def set_password(self, host, username, password):
    key = self.get_key()
    encrypt_secret = self.encrypt_string(key, password)
    sq = f"""update logins set password_value=x'{encrypt_secret.hex()}' where signon_realm ='{host}' and username_value='{username}';"""
    with sqlite3.connect(self.cookie_path) as conn:
      cu = conn.cursor()
      cu.execute(sq)
      conn.commit()


if __name__ == '__main__':
  a = Chrome()
  aa = a.get_password("https://baidu.com")
  print(aa)