from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import base64
import gradio as gr
from transformers import pipeline

# Modeli yükle
classifier = pipeline("text-classification", model="BaranKanat/BerTurk-SpamSMS")

# Sabit AES anahtarı (16, 24 veya 32 byte uzunlukta olmalı)
AES_KEY = b"ThisIsASecretKey"  # 16-byte (128-bit)

# PKCS7 padding
def pad(data: bytes) -> bytes:
    pad_len = 16 - len(data) % 16
    return data + bytes([pad_len] * pad_len)

def unpad(data: bytes) -> bytes:
    pad_len = data[-1]
    return data[:-pad_len]


# AES ile şifreleme (CBC mode), BASE64 çıktısı verir
def encrypt(text):
    if not text:
        raise ValueError("Şifrelenecek metin boş olamaz.")
    iv = get_random_bytes(16)
    cipher = AES.new(AES_KEY, AES.MODE_CBC, iv)
    data_bytes = text.encode('utf-8')  # Türkçe karakterleri de destekler
    padded_data = pad(data_bytes)
    encrypted_bytes = cipher.encrypt(padded_data)
    encrypted_message = base64.b64encode(iv + encrypted_bytes).decode('utf-8')
    return encrypted_message
    


# AES ile şifre çözme (BASE64 giriş)
def decrypt(encrypted_text):
    try:
        encrypted_data = base64.b64decode(encrypted_text)
        iv = encrypted_data[:16]
        encrypted_msg = encrypted_data[16:]
        cipher = AES.new(AES_KEY, AES.MODE_CBC, iv)
        decrypted_bytes = cipher.decrypt(encrypted_msg)
        unpadded_bytes = unpad(decrypted_bytes)
        return unpadded_bytes.decode('utf-8')
    except Exception as e:
        return f"Hata: {str(e)}"

# Şifreli SMS'i çözüp tahmin yapan fonksiyon
def predict_encrypted_sms(encrypted_sms):
    try:
        decrypted = decrypt(encrypted_sms)
        result = classifier(decrypted)[0]
        return f"🔓 Şifre Çözülmüş Metin:\n{decrypted}\n\n📊 Tahmin:\n{result['label']} (%{result['score']*100:.2f})"
    except Exception as e:
        return f"Hata oluştu: {str(e)}"

# Demo için metin şifreleme
def encrypt_for_demo(plain_sms):
    try:
        return encrypt(plain_sms)
    except Exception as e:
        return f"Hata oluştu: {str(e)}"

# Gradio arayüzü
demo = gr.TabbedInterface(
    interface_list=[
        gr.Interface(
            fn=predict_encrypted_sms,
            inputs=gr.Textbox(label="🔐 Şifreli SMS Metni Girin (Base64)"),
            outputs="text",
            title="BerTurk Spam SMS Tespiti (AES + Base64)",
            description="AES algoritmasıyla base64 biçiminde şifrelenmiş SMS metnini çözüp spam olup olmadığını tahmin eder."
        ),
        gr.Interface(
            fn=encrypt_for_demo,
            inputs=gr.Textbox(label="📨 Normal SMS Metni (Şifrelenmek Üzere)"),
            outputs="text",
            title="SMS Şifreleme (AES + Base64)",
            description="Buraya yazdığınız metni AES ile şifreler ve base64 biçiminde çıktı verir."
        )
    ],
    tab_names=["Tahmin (Şifre Çöz)", "Demo: Metin Şifrele"]
)

demo.launch()