app.py

from flask import Flask, render_template, request, redirect, url_for, session, jsonify
import psycopg2
import psycopg2.extras
import hashlib
import os
from functools import wraps
from dotenv import load_dotenv

# Load environment variables from .env file (for local development)
load_dotenv()

app = Flask(__name__)
app.secret_key = os.environ.get('SECRET_KEY', 'your-secret-key-change-in-production')

# ── ZTAM Zero-Trust integration ──────────────────────────────────────────────
# When this app runs behind the ZTAM Envoy gateway, the platform:
#   1. Validates the user's Keycloak token on EVERY request
#   2. Strips X-Username/X-User-Roles from incoming requests (prevents spoofing)
#   3. Re-injects X-Username and X-User-Roles only after successful validation
# This app just reads those two trusted headers. No JWT, no secrets needed.


@app.before_request
def ztam_auto_session():
    """Read identity from ZTAM gateway headers and establish a Flask session.

    Envoy strips X-Username / X-User-Roles from ALL incoming requests, then
    auth-middleware re-adds them only after Keycloak + OPA validation succeeds.
    """
    username = request.headers.get('X-Username', '').strip()
    roles    = request.headers.get('X-User-Roles', '').strip()

    if not username or not roles:
        return  # not behind ZTAM gateway — standalone login flow handles it

    # Avoid a DB round-trip on every request within the same browser session
    if session.get('username') == username:
        session['role'] = roles.split(',')[0]  # role may change between requests
        return

    # New user or session mismatch: look up the integer user_id from DB
    try:
        conn   = get_db()
        cursor = conn.cursor()
        cursor.execute('SELECT id FROM users WHERE username = %s', (username,))
        user = cursor.fetchone()
        cursor.close()
        conn.close()
        if user:
            session['user_id']  = user['id']
            session['username'] = username
            session['role']     = roles.split(',')[0]
    except Exception:
        pass  # DB unavailable — login_required will redirect to /login


def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

def get_db():
    """Get PostgreSQL database connection"""
    try:
        # Try DATABASE_URL first (for Render/Railway/production)
        database_url = os.environ.get('DATABASE_URL')
        if database_url:
            # Render gives postgres:// but psycopg2 needs postgresql://
            database_url = database_url.replace('postgres://', 'postgresql://', 1)
            conn = psycopg2.connect(database_url, cursor_factory=psycopg2.extras.RealDictCursor)
        else:
            # Fall back to individual environment variables (for local development)
            conn = psycopg2.connect(
                host=os.environ.get('DB_HOST', 'localhost'),
                user=os.environ.get('DB_USER', 'postgres'),
                password=os.environ.get('DB_PASSWORD', ''),
                database=os.environ.get('DB_NAME', 'store_db'),
                port=int(os.environ.get('DB_PORT', '5432')),
                cursor_factory=psycopg2.extras.RealDictCursor
            )
        return conn
    except Exception as e:
        print(f"Error connecting to PostgreSQL: {e}")
        raise

def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function

def role_required(roles):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if 'role' not in session or session['role'] not in roles:
                return "Access Denied", 403
            return f(*args, **kwargs)
        return decorated_function
    return decorator

@app.route('/')
def index():
    return redirect(url_for('login'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = hash_password(request.form['password'])
        
        conn = get_db()
        cursor = conn.cursor()
        cursor.execute('SELECT * FROM users WHERE username = %s AND password = %s', (username, password))
        user = cursor.fetchone()
        cursor.close()
        conn.close()
        
        if user:
            session['user_id'] = user['id']
            session['username'] = user['username']
            session['role'] = user['role']
            
            if user['role'] == 'admin':
                return redirect(url_for('admin'))
            elif user['role'] == 'manager':
                return redirect(url_for('manager'))
            else:
                return redirect(url_for('store'))
        else:
            return render_template('login.html', error='Invalid credentials')
    
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.clear()
    return redirect(url_for('login'))

@app.route('/store')
@login_required
def store():
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM products WHERE stock > 0')
    products = cursor.fetchall()
    cursor.close()
    conn.close()
    return render_template('store.html', products=products)

@app.route('/order', methods=['POST'])
@login_required
def order():
    product_id = request.form['product_id']
    quantity = int(request.form['quantity'])
    
    conn = get_db()
    cursor = conn.cursor()
    
    cursor.execute('SELECT * FROM products WHERE id = %s', (product_id,))
    product = cursor.fetchone()
    
    if product and product['stock'] >= quantity:
        total_price = product['price'] * quantity
        
        cursor.execute(
            'INSERT INTO orders (user_id, product_id, quantity, total_price) VALUES (%s, %s, %s, %s)',
            (session['user_id'], product_id, quantity, total_price)
        )
        
        cursor.execute(
            'UPDATE products SET stock = stock - %s WHERE id = %s',
            (quantity, product_id)
        )
        
        conn.commit()
        message = f"Order placed successfully! Total: ${total_price:.2f}"
    else:
        message = "Insufficient stock!"
    
    cursor.close()
    conn.close()
    
    return redirect(url_for('store'))

@app.route('/admin')
@login_required
@role_required(['admin'])
def admin():
    conn = get_db()
    cursor = conn.cursor()
    
    cursor.execute('SELECT * FROM users')
    users = cursor.fetchall()
    
    cursor.execute('SELECT * FROM products')
    products = cursor.fetchall()
    
    cursor.execute('''
        SELECT o.*, u.username, p.name as product_name 
        FROM orders o 
        JOIN users u ON o.user_id = u.id 
        JOIN products p ON o.product_id = p.id 
        ORDER BY o.order_date DESC
    ''')
    orders = cursor.fetchall()
    
    cursor.close()
    conn.close()
    
    return render_template('admin.html', users=users, products=products, orders=orders)

@app.route('/manager')
@login_required
@role_required(['manager', 'admin'])
def manager():
    conn = get_db()
    cursor = conn.cursor()
    
    cursor.execute('SELECT * FROM products')
    products = cursor.fetchall()
    
    cursor.execute('''
        SELECT o.*, u.username, p.name as product_name 
        FROM orders o 
        JOIN users u ON o.user_id = u.id 
        JOIN products p ON o.product_id = p.id 
        ORDER BY o.order_date DESC 
        LIMIT 50
    ''')
    orders = cursor.fetchall()
    
    cursor.close()
    conn.close()
    
    return render_template('manager.html', products=products, orders=orders)

@app.route('/add_product', methods=['POST'])
@login_required
@role_required(['manager', 'admin'])
def add_product():
    name = request.form['name']
    description = request.form['description']
    price = float(request.form['price'])
    stock = int(request.form['stock'])
    
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO products (name, description, price, stock) VALUES (%s, %s, %s, %s)',
        (name, description, price, stock)
    )
    conn.commit()
    cursor.close()
    conn.close()
    
    if session['role'] == 'admin':
        return redirect(url_for('admin'))
    return redirect(url_for('manager'))

@app.route('/update_stock', methods=['POST'])
@login_required
@role_required(['manager', 'admin'])
def update_stock():
    product_id = request.form['product_id']
    new_stock = int(request.form['stock'])
    
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('UPDATE products SET stock = %s WHERE id = %s', (new_stock, product_id))
    conn.commit()
    cursor.close()
    conn.close()
    
    if session['role'] == 'admin':
        return redirect(url_for('admin'))
    return redirect(url_for('manager'))

# REST API Endpoints for testing
@app.route('/api/test/health')
def health():
    return jsonify({"status": "healthy", "database": "postgresql"})

@app.route('/api/products')
def api_products():
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM products')
    products = cursor.fetchall()
    cursor.close()
    conn.close()
    return jsonify(products)

@app.route('/api/users')
@login_required
@role_required(['admin'])
def api_users():
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('SELECT id, username, role, email FROM users')
    users = cursor.fetchall()
    cursor.close()
    conn.close()
    return jsonify(users)

@app.route('/api/orders')
@login_required
def api_orders():
    conn = get_db()
    cursor = conn.cursor()
    
    if session['role'] == 'admin':
        cursor.execute('''
            SELECT o.*, u.username, p.name as product_name 
            FROM orders o 
            JOIN users u ON o.user_id = u.id 
            JOIN products p ON o.product_id = p.id 
            ORDER BY o.order_date DESC
        ''')
    else:
        cursor.execute('''
            SELECT o.*, p.name as product_name 
            FROM orders o 
            JOIN products p ON o.product_id = p.id 
            WHERE o.user_id = %s 
            ORDER BY o.order_date DESC
        ''', (session['user_id'],))
    
    orders = cursor.fetchall()
    cursor.close()
    conn.close()
    return jsonify(orders)

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)