import os
import io
import struct
import base64
from flask import Flask, request, render_template_string, send_file
from PIL import Image
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

app = Flask(__name__)
Image.MAX_IMAGE_PIXELS = None

# --- Core Logic (The Vault) ---

def derive_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    return kdf.derive(password.encode())

def create_vault(file_bytes, filename, password=None):
    # Prepare payload: [name_len (1) | name | file_len (4) | file_data]
    name_enc = filename.encode('utf-8')
    payload = struct.pack('B', len(name_enc)) + name_enc + struct.pack('<I', len(file_bytes)) + file_bytes
    
    if password:
        salt = os.urandom(16)
        iv = os.urandom(12)
        key = derive_key(password, salt)
        aesgcm = AESGCM(key)
        encrypted_data = aesgcm.encrypt(iv, payload, None)
        # Format: [version (1) | salt (16) | iv (12) | data_len (4) | data]
        final_blob = b'\x01' + salt + iv + struct.pack('<I', len(encrypted_data)) + encrypted_data
    else:
        # Format: [version (1) | payload]
        final_blob = b'\x00' + payload

    # Calculate image dimensions (3 bytes per pixel)
    num_pixels = (len(final_blob) + 2) // 3
    side = max(16, int(num_pixels**0.5) + 1)
    
    # Pad blob to fit pixels exactly
    padded_blob = final_blob.ljust(side * side * 3, b'\x00')
    
    # Create Image
    img = Image.frombytes('RGB', (side, side), padded_blob)
    img_io = io.BytesIO()
    img.save(img_io, 'PNG')
    img_io.seek(0)
    return img_io, "vault_secure.png" if password else "vault_public.png"

def extract_vault(img_bytes, password=None):
    img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
    flat_data = img.tobytes()
    
    version = flat_data[0]
    
    if version == 1: # Encrypted
        if not password: raise ValueError("Password required for this vault.")
        salt = flat_data[1:17]
        iv = flat_data[17:29]
        data_len = struct.unpack('<I', flat_data[29:33])[0]
        encrypted_data = flat_data[33:33+data_len]
        
        key = derive_key(password, salt)
        aesgcm = AESGCM(key)
        decrypted_payload = aesgcm.decrypt(iv, encrypted_data, None)
    else: # Public
        decrypted_payload = flat_data[1:]

    # Parse Payload
    name_len = decrypted_payload[0]
    filename = decrypted_payload[1:1+name_len].decode('utf-8')
    file_size = struct.unpack('<I', decrypted_payload[1+name_len:1+name_len+4])[0]
    actual_file_data = decrypted_payload[1+name_len+4 : 1+name_len+4+file_size]
    
    return io.BytesIO(actual_file_data), filename

# --- Web UI ---

HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <title>PNG Vault Python</title>
    <style>
        body { font-family: sans-serif; background: #0B0F1A; color: white; display: flex; justify-content: center; padding: 50px; }
        .card { background: #161C2D; padding: 30px; border-radius: 20px; width: 400px; border: 1px solid #2DD4BF; }
        input, button { width: 100%; margin: 10px 0; padding: 10px; border-radius: 5px; border: none; box-sizing: border-box; }
        button { background: #2DD4BF; color: #064e3b; font-weight: bold; cursor: pointer; }
        .secondary { background: #334155; color: white; }
        h2 { text-align: center; color: #2DD4BF; }
    </style>
</head>
<body>
    <div class="card">
        <h2>PNG Vault</h2>
        <form action="/encode" method="post" enctype="multipart/form-data">
            <small>HIDE FILE</small>
            <input type="file" name="file" required>
            <input type="password" name="password" placeholder="Optional Password">
            <button type="submit">Create Vault Image</button>
        </form>
        <hr style="opacity:0.1; margin: 20px 0;">
        <form action="/decode" method="post" enctype="multipart/form-data">
            <small>EXTRACT FILE</small>
            <input type="file" name="file" accept="image/png" required>
            <input type="password" name="password" placeholder="Password (if encrypted)">
            <button type="submit" class="secondary">Extract Hidden File</button>
        </form>
    </div>
</body>
</html>
"""

@app.route('/')
def index():
    return render_template_string(HTML_TEMPLATE)

@app.route('/encode', methods=['POST'])
def encode_route():
    file = request.files['file']
    password = request.form.get('password')
    img_io, name = create_vault(file.read(), file.filename, password)
    return send_file(img_io, mimetype='image/png', as_attachment=True, download_name=name)

@app.route('/decode', methods=['POST'])
def decode_route():
    file = request.files['file']
    password = request.form.get('password')
    try:
        file_io, name = extract_vault(file.read(), password)
        return send_file(file_io, mimetype='application/octet-stream', as_attachment=True, download_name=name)
    except Exception as e:
        return f"Error: {str(e)}", 400

if __name__ == '__main__':
    app.run(debug=True, port=8080)
