ㅁㅁㅁ

import tkinter as tk
from tkinter import ttk, scrolledtext, simpledialog
import requests
from Cryptodome.Cipher import AES
from Cryptodome.Util.Padding import pad, unpad
import base64
import threading
import time
import random
import string
import hashlib

class ConsoleStyle:
    BG_COLOR = 'black'
    FG_COLOR = '#00ff00'  # 연한 녹색 (터미널 스타일)
    FONT = ('Consolas', 10)  # 고정폭 폰트
    CURSOR_COLOR = '#00ff00'
    SELECT_BG = '#003300'  # 선택 영역 배경색
    SELECT_FG = '#00ff00'  # 선택 영역 텍스트색

class ChatClientGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Secure Terminal Chat")
        self.setup_gui()
        self.setup_client()
        self.apply_console_style()
        self.setup_shortcuts()  # 단축키 설정 추가
        
        # 암호화 관련 변수 초기화
        self.secret_key = None
        self.iv = None

    def setup_shortcuts(self):
        # Ctrl+O 단축키 바인딩
        self.root.bind('<Control-o>', lambda e: self.clear_chat())

    def clear_chat(self):
        # 채팅 영역의 모든 내용을 삭제
        self.chat_area.delete(1.0, tk.END)
        self.show_system_message("대화 내용이 삭제되었습니다.")

    def apply_console_style(self):
        # 윈도우 스타일
        self.root.configure(bg=ConsoleStyle.BG_COLOR)
        
        # 스타일 설정
        style = ttk.Style()
        style.configure('Console.TFrame', background=ConsoleStyle.BG_COLOR)
        style.configure('Console.TLabel', 
                       background=ConsoleStyle.BG_COLOR,
                       foreground=ConsoleStyle.FG_COLOR,
                       font=ConsoleStyle.FONT)
        style.configure('Console.TButton',
                       background=ConsoleStyle.BG_COLOR,
                       foreground=ConsoleStyle.FG_COLOR,
                       font=ConsoleStyle.FONT)
        style.configure('Console.TEntry',
                       background=ConsoleStyle.BG_COLOR,
                       foreground=ConsoleStyle.FG_COLOR,
                       font=ConsoleStyle.FONT,
                       fieldbackground=ConsoleStyle.BG_COLOR)

        # 채팅 영역 스타일
        self.chat_area.configure(
            bg=ConsoleStyle.BG_COLOR,
            fg=ConsoleStyle.FG_COLOR,
            font=ConsoleStyle.FONT,
            insertbackground=ConsoleStyle.CURSOR_COLOR,
            selectbackground=ConsoleStyle.SELECT_BG,
            selectforeground=ConsoleStyle.SELECT_FG
        )

        # 입력 영역 스타일
        self.message_entry.configure(
            style='Console.TEntry',
            font=ConsoleStyle.FONT
        )

    def setup_client(self):
        self.client_id = self.generate_client_id()
        self.running = True
        self.server_url = None
        
        # 서버 연결 프레임
        self.connect_frame = ttk.Frame(self.root, padding="10", style='Console.TFrame')
        self.connect_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        ttk.Label(self.connect_frame, text="서버 IP:", style='Console.TLabel').grid(row=0, column=0, padx=5)
        self.server_ip = ttk.Entry(self.connect_frame, width=20, style='Console.TEntry')
        self.server_ip.insert(0, "localhost")
        self.server_ip.grid(row=0, column=1, padx=5)
        
        # PIN 입력 추가
        ttk.Label(self.connect_frame, text="PIN (4자리):", style='Console.TLabel').grid(row=1, column=0, padx=5, pady=5)
        self.pin_entry = ttk.Entry(self.connect_frame, width=10, style='Console.TEntry', show="*")
        self.pin_entry.grid(row=1, column=1, padx=5, pady=5, sticky=tk.W)
        
        self.connect_btn = ttk.Button(self.connect_frame, text="연결", 
                                    command=self.connect_to_server, 
                                    style='Console.TButton')
        self.connect_btn.grid(row=1, column=2, padx=5)

        # 메시지 영역은 처음에는 숨김
        self.message_frame.grid_remove()

    def setup_gui(self):
        # 메시지 프레임
        self.message_frame = ttk.Frame(self.root, padding="10", style='Console.TFrame')
        self.message_frame.grid(row=1, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 채팅 영역
        self.chat_area = scrolledtext.ScrolledText(
            self.message_frame, 
            wrap=tk.WORD, 
            width=80,  # 더 넓게
            height=24  # 더 높게
        )
        self.chat_area.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 메시지 입력
        self.message_entry = ttk.Entry(
            self.message_frame,
            style='Console.TEntry',
            width=70  # 더 넓게
        )
        self.message_entry.grid(row=1, column=0, sticky=(tk.W, tk.E), padx=5, pady=5)
        self.message_entry.bind('<Return>', lambda e: self.send_message())
        
        # 전송 버튼
        self.send_button = ttk.Button(
            self.message_frame, 
            text="전송", 
            command=self.send_message,
            style='Console.TButton'
        )
        self.send_button.grid(row=1, column=1, sticky=(tk.E), pady=5)
        
        # Grid 설정
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(1, weight=1)
        self.message_frame.columnconfigure(0, weight=1)

    def generate_client_id(self):
        return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
    
    def generate_key_from_pin(self, pin):
        """PIN을 32바이트 키와 16바이트 IV로 변환"""
        # PIN을 바이트로 변환
        pin_bytes = pin.encode('utf-8')
        
        # SHA-256으로 해시하여 32바이트 키 생성
        key_hash = hashlib.sha256(pin_bytes).digest()
        
        # MD5로 해시하여 16바이트 IV 생성
        iv_hash = hashlib.md5(pin_bytes).digest()
        
        return key_hash, iv_hash

    def connect_to_server(self):
        # PIN 입력 확인
        pin = self.pin_entry.get().strip()
        if not pin or len(pin) != 4 or not pin.isdigit():
            self.show_system_message("유효한 4자리 PIN을 입력해주세요.")
            return
            
        # PIN에서 암호화 키와 IV 생성
        self.secret_key, self.iv = self.generate_key_from_pin(pin)
        
        server_ip = self.server_ip.get()
        self.server_url = f'http://{server_ip}:5000'
        
        try:
            response = requests.post(f"{self.server_url}/register", 
                                   json={"client_id": self.client_id})
            if response.status_code == 200:
                self.show_system_message(f"서버에 연결되었습니다. (PIN: {pin})")
                self.connect_frame.grid_remove()
                self.message_frame.grid()
                self.receive_thread = threading.Thread(target=self.receive_messages)
                self.receive_thread.daemon = True
                self.receive_thread.start()
            else:
                self.show_system_message("서버 연결 실패")
        except Exception as e:
            self.show_system_message(f"연결 오류: {str(e)}")

    def encrypt_message(self, message):
        if not self.secret_key or not self.iv:
            self.show_system_message("암호화 키가 설정되지 않았습니다.")
            return None
            
        try:
            message_bytes = message.encode('utf-8')
            padded_message = pad(message_bytes, AES.block_size)
            cipher = AES.new(self.secret_key, AES.MODE_CBC, IV=self.iv)
            encrypted_message = cipher.encrypt(padded_message)
            return base64.b64encode(encrypted_message).decode('utf-8')
        except Exception as e:
            self.show_system_message(f"암호화 오류: {str(e)}")
            return None

    def decrypt_message(self, encrypted_message):
        if not self.secret_key or not self.iv:
            self.show_system_message("복호화 키가 설정되지 않았습니다.")
            return None
            
        try:
            encrypted_bytes = base64.b64decode(encrypted_message)
            cipher = AES.new(self.secret_key, AES.MODE_CBC, IV=self.iv)
            decrypted_padded = cipher.decrypt(encrypted_bytes)
            return unpad(decrypted_padded, AES.block_size).decode('utf-8')
        except Exception as e:
            # 복호화 실패 시 상세 오류는 콘솔에만 출력
            print(f"복호화 오류: {str(e)}")
            return "[암호화된 메시지 - PIN이 일치하지 않습니다]"

    def send_message(self):
        if not self.server_url:
            return
            
        message = self.message_entry.get().strip()
        if message:
            encrypted_message = self.encrypt_message(message)
            if encrypted_message:
                try:
                    response = requests.post(
                        f"{self.server_url}/send",
                        json={
                            "sender_id": self.client_id,
                            "message": encrypted_message
                        }
                    )
                    if response.status_code == 200:
                        self.message_entry.delete(0, tk.END)
                        self.show_message("me", message)
                except Exception as e:
                    self.show_system_message(f"전송 오류: {str(e)}")

    def receive_messages(self):
        while self.running:
            if not self.server_url:
                time.sleep(1)
                continue
                
            try:
                response = requests.post(
                    f"{self.server_url}/receive",
                    json={"client_id": self.client_id}
                )
                if response.status_code == 200:
                    messages = response.json()['messages']
                    for encrypted_message in messages:
                        try:
                            decrypted_message = self.decrypt_message(encrypted_message)
                            if decrypted_message:
                                self.root.after(0, self.show_message, "other", decrypted_message)
                        except Exception as e:
                            print(f"메시지 처리 오류: {str(e)}")
            except Exception as e:
                print(f"수신 오류: {str(e)}")
            time.sleep(1)

    def show_message(self, sender, message):
        timestamp = time.strftime("%H:%M:%S")
        if sender == "me":
            formatted_message = f"[{timestamp}] >>> {message}\n"
        else:
            # 상대방 메시지는 ID 없이 표시
            formatted_message = f"[{timestamp}] <<< {message}\n"
        self.chat_area.insert(tk.END, formatted_message)
        self.chat_area.see(tk.END)

    def show_system_message(self, message):
        timestamp = time.strftime("%H:%M:%S")
        formatted_message = f"[{timestamp}] === {message} ===\n"
        self.chat_area.insert(tk.END, formatted_message)
        self.chat_area.see(tk.END)

    def on_closing(self):
        self.running = False
        if self.server_url:
            try:
                requests.post(f"{self.server_url}/unregister", 
                            json={"client_id": self.client_id})
            except:
                pass
        self.root.destroy()

def main():
    root = tk.Tk()
    root.configure(bg='black')
    # 윈도우 크기 설정
    root.geometry("800x600")  # 더 큰 창 크기
    app = ChatClientGUI(root)
    root.protocol("WM_DELETE_WINDOW", app.on_closing)
    root.mainloop()

if __name__ == '__main__':
    main()