app.py

from flask import Flask, render_template, redirect, url_for, request, flash, jsonify, Response
import os
import json
import datetime
import threading
import time
import requests
from wakeonlan import send_magic_packet
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
from ad_integration import authenticate_user, sync_users_from_ad

app = Flask(__name__)
app.secret_key = 'your_secret_key'

# Lock to prevent concurrent read-modify-write on workstations.json
workstation_lock = threading.Lock()

def time_ago(timestamp_str):
    if not timestamp_str:
        return ""
    try:
        past_time = datetime.datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
        now = datetime.datetime.now()
        diff = now - past_time
        minutes = int(diff.total_seconds() / 60)
        
        if minutes < 1:
            return "Just now"
        elif minutes == 1:
            return "1 min ago"
        elif minutes < 60:
            return f"{minutes} mins ago"
        else:
            hours = minutes // 60
            if hours < 24:
                return f"{hours} hours ago"
            else:
                days = hours // 24
                return f"{days} days ago"
    except Exception:
        return timestamp_str

app.jinja_env.filters['time_ago'] = time_ago

def parse_datetime(timestamp_str):
    """Parse a datetime string into a datetime object."""
    if not timestamp_str:
        return datetime.datetime.min
    try:
        return datetime.datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
    except Exception:
        return datetime.datetime.min

app.jinja_env.filters['parse_datetime'] = parse_datetime

login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'

class User(UserMixin):
    def __init__(self, username, data):
        self.id = username
        self.email = data.get('email', '')
        self.assigned_macs = data.get('assigned_macs', [])
        self.is_admin = data.get('is_admin', False)

@login_manager.user_loader
def load_user(username):
    users = read_users()
    if username in users:
        return User(username, users[username])
    return None

def load_json(file_path, default=None):
    if os.path.exists(file_path):
        try:
            with open(file_path, 'r', encoding='utf-8-sig') as file:
                return json.load(file)
        except json.JSONDecodeError as e:
            print(f"DTO: CRITICAL - Error decoding {file_path}: {e}")
            # Do NOT return empty list/dict on decode error — that causes data wipes
            # when a subsequent write saves the empty data back to disk.
            raise
    if default is not None:
        return default
    return [] if 'workstations' in file_path else {}

def save_json(file_path, data):
    with open(file_path, 'w') as file:
        json.dump(data, file, indent=4)

def read_workstations():
    return load_json('workstations.json')

def write_workstation(name, ip, mac):
    with workstation_lock:
        workstations = read_workstations()
        workstations.append({'name': name, 'ip': ip, 'mac': mac})
        save_json('workstations.json', workstations)

def delete_workstation(mac):
    with workstation_lock:
        workstations = read_workstations()
        workstations = [ws for ws in workstations if ws['mac'] != mac]
        save_json('workstations.json', workstations)

def read_users():
    return load_json('users.json')

def save_users(users):
    save_json('users.json', users)

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        # In production, use authenticate_user(username, password)
        # For testing without AD, we can have a bypass or mock
        if authenticate_user(username, password) or (username == 'admin' and password == 'admin'):
            users = read_users()
            if username not in users:
                # Auto-register valid AD user
                users[username] = {'email': '', 'assigned_macs': [], 'is_admin': False}
                save_users(users)
            
            user = User(username, users[username])
            login_user(user)
            return redirect(url_for('home'))
        else:
            flash('Invalid credentials', 'danger')
            
    return render_template('login.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))

@app.route('/')
@login_required
def home():
    all_workstations = read_workstations()
    # User requested that even admins should only see assigned computers on the home page
    workstations = [ws for ws in all_workstations if ws['mac'] in current_user.assigned_macs]
    return render_template('index.html', workstations=workstations, user=current_user, now=datetime.datetime.now())

@app.route('/wake/<ip>/<mac>')
def wake(ip, mac):
    send_magic_packet(mac, ip_address=ip)
    flash('Magic packet sent successfully!', 'success')
    return redirect(url_for('home'))


@app.route('/add', methods=['POST'])
@login_required
def add():
    if not current_user.is_admin:
        flash('Access denied.', 'danger')
        return redirect(url_for('home'))
        
    name = request.form['name']
    ip = request.form['ip']
    mac = request.form['mac']
    write_workstation(name, ip, mac)
    flash('Workstation added successfully!', 'success')
    # Redirect back to admin since this action is now there
    return redirect(url_for('admin'))

@app.route('/delete', methods=['POST'])
@login_required
def delete():
    if not current_user.is_admin:
        flash('Access denied.', 'danger')
        return redirect(url_for('home'))
        
    mac = request.form['mac']
    delete_workstation(mac)
    flash('Workstation deleted successfully!', 'success')
     # Redirect back to admin since this action is now there
    return redirect(url_for('admin'))

@app.route('/admin')
@login_required
def admin():
    if not current_user.is_admin:
        flash('Access denied.', 'danger')
        return redirect(url_for('home'))
    
    users = read_users()
    workstations = read_workstations()
    return render_template('admin.html', users=users, workstations=workstations)

@app.route('/admin/toggle_admin', methods=['POST'])
@login_required
def toggle_admin():
    if not current_user.is_admin:
        return redirect(url_for('home'))
        
    username = request.form['username']
    users = read_users()
    if username in users:
        users[username]['is_admin'] = not users[username].get('is_admin', False)
        save_users(users)
        flash(f"Admin status for {username} changed.", 'success')
    return redirect(url_for('admin'))

@app.route('/admin/assign', methods=['POST'])
@login_required
def assign_workstations():
    if not current_user.is_admin:
        return redirect(url_for('home'))
        
    username = request.form['username']
    assigned_macs = request.form.getlist('assigned_macs')
    
    users = read_users()
    if username in users:
        users[username]['assigned_macs'] = assigned_macs
        save_users(users)
        flash(f"Assignments updated for {username}.", 'success')
    return redirect(url_for('admin'))

@app.route('/admin/sync', methods=['POST'])
@login_required
def sync_users():
    if not current_user.is_admin:
        return redirect(url_for('home'))
        
    synced_users = sync_users_from_ad({}) # Pass empty or existing users? Function signature expects current_users? 
    # Checking ad_integration.py: def sync_users_from_ad(current_users):
    
    current_users = read_users()
    new_users = sync_users_from_ad(current_users)
    
    # Merge logic: don't overwrite existing assignments/admin status
    for username, data in new_users.items():
        if username not in current_users:
            current_users[username] = {
                'email': data.get('email', ''),
                'assigned_macs': [],
                'is_admin': False
            }
        else:
            # Update email if changed, keep others
            current_users[username]['email'] = data.get('email', current_users[username].get('email', ''))
            
        if username == 'gabor.abbas':
             current_users[username]['is_admin'] = True

    save_users(current_users)
    save_json('last_sync.json', {'last_sync': datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")})
    flash('Users synced from Active Directory.', 'success')
    return redirect(url_for('admin'))

@app.route('/api/register', methods=['POST'])
def register_workstation():
    data = request.get_json()
    if not data:
        return jsonify({'error': 'Invalid JSON'}), 400
    
    mac = data.get('mac')
    ip = data.get('ip')
    name = data.get('name')
    last_seen = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    if not mac or not ip or not name:
        return jsonify({'error': 'Missing required fields'}), 400
    
    # ── Handle both old and new payload formats ──
    active_users = data.get('active_users', None)
    
    if active_users is None:
        # Old format: single "user" + "idle_seconds" fields (backwards compat)
        old_user = data.get('user', '')
        old_idle = data.get('idle_seconds', 0)
        if old_user:
            active_users = [{"username": old_user, "idle_seconds": old_idle, "session_type": "unknown"}]
        else:
            active_users = []
    
    # Derive last_user and idle_seconds for backwards compat with UI
    if active_users:
        # Prefer console user, fall back to first user
        console_users = [u for u in active_users if u.get('session_type') == 'console']
        primary_user = console_users[0] if console_users else active_users[0]
        last_user = primary_user['username']
        # Use minimum idle across all users (most recently active)
        idle_seconds = min(u.get('idle_seconds', 0) for u in active_users)
    else:
        last_user = ''
        idle_seconds = 0
    
    with workstation_lock:
        try:
            workstations = load_json('workstations.json')
        except json.JSONDecodeError:
            return jsonify({'error': 'Server data file corrupted, skipping write to prevent data loss'}), 500
        
        # Check if workstation exists
        existing = next((item for item in workstations if item['mac'] == mac), None)
        
        # ── Track disconnected_since per user (server-side) ──
        # Build a lookup of previously stored disconnect timestamps
        prev_disconnect_times = {}
        if existing and existing.get('active_users'):
            for prev_u in existing['active_users']:
                if prev_u.get('disconnected_since'):
                    prev_disconnect_times[prev_u['username']] = prev_u['disconnected_since']
        
        # Stamp disconnected_since on disconnected users
        for u in active_users:
            if u.get('session_type') == 'disconnected':
                # Keep existing timestamp if user was already disconnected, else stamp now
                u['disconnected_since'] = prev_disconnect_times.get(u['username'], last_seen)
            else:
                # Active users - no disconnect timestamp
                u.pop('disconnected_since', None)
        
        if existing:
            existing['ip'] = ip
            existing['name'] = name
            existing['last_user'] = last_user
            existing['last_seen'] = last_seen
            existing['idle_seconds'] = idle_seconds
            existing['active_users'] = active_users
        else:
            workstations.append({
                'mac': mac,
                'ip': ip,
                'name': name,
                'last_user': last_user,
                'last_seen': last_seen,
                'idle_seconds': idle_seconds,
                'active_users': active_users
            })
            
        save_json('workstations.json', workstations)
    return jsonify({'status': 'success', 'message': 'Workstation registered'}), 200

@app.route('/rdp/<ip>')
@login_required
def rdp_download(ip):
    # RDP file content
    rdp_content = f"full address:s:{ip}\nprompt for credentials:i:1"
    
    return Response(
        rdp_content,
        mimetype="application/x-rdp",
        headers={"Content-disposition": f"attachment; filename={ip}.rdp"}
    )

def reset_plug_thread(plug_id, ha_url, ha_token):
    headers = {
        "Authorization": f"Bearer {ha_token}",
        "Content-Type": "application/json",
    }
    
    # Turn OFF
    try:
        print(f"DTO: Turning OFF plug {plug_id} via {ha_url}...")
        resp = requests.post(f"{ha_url}/api/services/switch/turn_off", headers=headers, json={"entity_id": plug_id})
        print(f"DTO: HA Response ({resp.status_code}): {resp.text}")
    except Exception as e:
        print(f"Error turning off plug {plug_id}: {e}")
        return

    # Wait 15 seconds
    time.sleep(15)

    # Turn ON
    try:
        print(f"DTO: Turning ON plug {plug_id}...")
        resp = requests.post(f"{ha_url}/api/services/switch/turn_on", headers=headers, json={"entity_id": plug_id})
        print(f"DTO: HA Response ({resp.status_code}): {resp.text}")
    except Exception as e:
        print(f"Error turning on plug {plug_id}: {e}")

@app.route('/hard_reset/<mac>', methods=['POST'])
@login_required
def hard_reset(mac):
    try:
        # Load configs
        plugs = load_json('plugs.json', {})
        ha_config = load_json('ha_config.json', {})
        workstations = load_json('workstations.json')
        
        # Find workstation by MAC to get Hostname
        target_ws = next((ws for ws in workstations if ws['mac'] == mac), None)
        
        if not target_ws:
            return jsonify({'status': 'error', 'message': 'Workstation not found'}), 404
            
        ws_name = target_ws.get('name')
        if not ws_name:
             return jsonify({'status': 'error', 'message': 'Workstation has no name'}), 400

        # Lookup plug by Hostname
        plug_id = plugs.get(ws_name)
        if not plug_id:
            # Auto-generate based on convention: PBV-LEVI -> switch.pbv_levi_switch_0
            sanitized_name = ws_name.lower().replace('-', '_')
            plug_id = f"switch.{sanitized_name}_switch_0"
            print(f"DTO: Auto-generated plug ID for {ws_name}: {plug_id}")
        
        # SAFETY GUARD: Prevent resetting the host machine
        if ws_name == "PBV-Mufasa":
            print("SAFETY BLOCK: Attempted to reset PBV-Mufasa. Aborting.")
            return jsonify({'status': 'error', 'message': 'Safety Lock: Cannot reset Host Machine (PBV-Mufasa)'}), 403
        
        ha_url = ha_config.get('url')
        ha_token = ha_config.get('token')
        
        if not ha_url or not ha_token:
            print("Warning: HA Config using defaults or missing.")

        # Start background thread
        thread = threading.Thread(target=reset_plug_thread, args=(plug_id, ha_url, ha_token))
        thread.start()
        
        return jsonify({'status': 'success', 'message': f'Hard reset initiated for {ws_name}'}), 200
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({'status': 'error', 'message': str(e)}), 500

@app.route('/api/power_status/<name>')
@login_required
def get_power_status(name):
    """Get power consumption status for a workstation's smart plug."""
    try:
        ha_config = load_json('ha_config.json', {})
        ha_url = ha_config.get('url')
        ha_token = ha_config.get('token')
        
        if not ha_url or not ha_token:
            return jsonify({'status': 'error', 'message': 'Home Assistant not configured'}), 500
        
        # Generate sensor entity ID based on naming convention
        # PBV-LEVI -> sensor.pbv_levi_switch_0_power
        sanitized_name = name.lower().replace('-', '_')
        sensor_id = f"sensor.{sanitized_name}_switch_0_power"
        
        headers = {
            "Authorization": f"Bearer {ha_token}",
            "Content-Type": "application/json",
        }
        
        resp = requests.get(f"{ha_url}/api/states/{sensor_id}", headers=headers)
        
        if resp.status_code == 200:
            data = resp.json()
            state = data.get('state', '0')
            try:
                power_watts = float(state)
            except (ValueError, TypeError):
                power_watts = 0
            
            is_running = power_watts > 70
            return jsonify({
                'status': 'success',
                'name': name,
                'power_watts': power_watts,
                'is_running': is_running
            }), 200
        else:
            return jsonify({
                'status': 'success',
                'name': name,
                'power_watts': 0,
                'is_running': False,
                'plug_not_found': True
            }), 200
            
    except Exception as e:
        print(f"Error getting power status for {name}: {e}")
        return jsonify({'status': 'error', 'message': str(e)}), 500

def check_ad_alive():
    try:
        import socket
        ad_config = load_json('ad_config.json', {})
        ip = ad_config.get('AD_SERVER_IP')
        port = ad_config.get('AD_SERVER_PORT', 389)
        if not ip:
            return False
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(2)
        s.connect((ip, int(port)))
        s.close()
        return True
    except Exception:
        return False

def check_ha_alive():
    try:
        ha_config = load_json('ha_config.json', {})
        url = ha_config.get('url')
        token = ha_config.get('token')
        if not url:
            return False
        headers = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"
        resp = requests.get(f"{url}/api/", headers=headers, timeout=2)
        return resp.status_code in [200, 401]
    except Exception:
        return False

def check_deadline_alive():
    try:
        deadline_config = load_json('deadline_config.json', {})
        url = deadline_config.get('url')
        if not url:
            return False
        resp = requests.get(f"{url}/api/slaves?NamesOnly=true", timeout=2)
        return resp.status_code == 200
    except Exception:
        return False

@app.route('/api/health')
@login_required
def get_health_status():
    """Check connectivity to Active Directory, Home Assistant, and Deadline."""
    ad_alive = check_ad_alive()
    ha_alive = check_ha_alive()
    deadline_alive = check_deadline_alive()
    
    last_sync_data = load_json('last_sync.json', {})
    last_sync = last_sync_data.get('last_sync', 'Never')
    
    return jsonify({
        'status': 'success',
        'ad': {
            'status': 'online' if ad_alive else 'offline',
            'last_sync': last_sync
        },
        'ha': {
            'status': 'online' if ha_alive else 'offline'
        },
        'deadline': {
            'status': 'online' if deadline_alive else 'offline'
        }
    }), 200

@app.route('/api/deadline_status')
@login_required
def get_deadline_status():
    """Query Thinkbox Deadline Web Service for worker statuses."""
    try:
        deadline_config = load_json('deadline_config.json', {})
        deadline_url = deadline_config.get('url')
        if not deadline_url:
            return jsonify({'status': 'error', 'message': 'Deadline config or URL is missing'}), 400

        # Fetch from the Deadline REST API
        # By default, Deadline Web Service listens on http://<server>:8081/api/slaves
        resp = requests.get(f"{deadline_url}/api/slaves", timeout=5)
        if resp.status_code == 200:
            workers_data = resp.json()
            
            # Map lowercase short hostname to its state
            worker_states = {}
            for w in workers_data:
                info = w.get('Info', {})
                settings = w.get('Settings', {})
                
                # Fetch Name from Info (fallback to Settings)
                name = info.get('Name') or settings.get('Name')
                
                # Determine state: Enabled = False in settings overrides Stat
                if settings.get('Enable') is False:
                    state = 'Disabled'
                else:
                    stat_code = info.get('Stat')
                    if stat_code == 1:
                        state = 'Rendering'
                    elif stat_code == 2:
                        state = 'Idle'
                    elif stat_code == 3:
                        state = 'Offline'
                    elif stat_code == 4:
                        state = 'Stalled'
                    else:
                        state = 'Offline'
                
                if name:
                    # Strip domain suffix (e.g. "pbv-foeldy.friday.local" -> "pbv-foeldy")
                    short_name = name.split('.')[0].lower()
                    worker_states[short_name] = state

            return jsonify({
                'status': 'success',
                'workers': worker_states
            }), 200
        else:
            return jsonify({
                'status': 'error',
                'message': f'Deadline service returned status code {resp.status_code}'
            }), resp.status_code
    except requests.exceptions.RequestException as e:
        # Gracefully handle connection errors so the page works even if Deadline is down
        print(f"Error querying Deadline Web Service: {e}")
        return jsonify({
            'status': 'error',
            'message': f'Could not connect to Deadline Web Service: {str(e)}'
        }), 502
    except Exception as e:
        print(f"Unexpected error in get_deadline_status: {e}")
        return jsonify({'status': 'error', 'message': str(e)}), 500

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