#!/usr/bin/env python3
"""
WebSocket SSH Daemon (wssshd)
Handles WebSocket connections from clients and wsssh/wsscp applications.

Copyright (C) 2024 Stefy Lanza <stefy@nexlab.net> and SexHack.me

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import argparse
import asyncio
import ssl
import websockets
import json
import sys
import os
import threading
import uuid
import subprocess
import pty
import select
import fcntl
import termios
import stat
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, send_from_directory
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash

# Client registry: id -> websocket
clients = {}
# Active tunnels: request_id -> {'client_ws': ws, 'wsssh_ws': ws, 'client_id': id}
active_tunnels = {}
# Active terminals: request_id -> {'client_id': id, 'username': username, 'proc': proc}
active_terminals = {}
debug = False
server_password = None
args = None

# Flask app for web interface
app = Flask(__name__)
app.config['SECRET_KEY'] = 'wsssh-secret-key-change-in-production'
config_dir = os.path.expanduser('~/.config/wssshd')
os.makedirs(config_dir, exist_ok=True)
app.config['SQLALCHEMY_DATABASE_URI'] = f'sqlite:///{config_dir}/users.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

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

class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(150), unique=True, nullable=False)
    password_hash = db.Column(db.String(150), nullable=False)
    is_admin = db.Column(db.Boolean, default=False)

@login_manager.user_loader
def load_user(user_id):
    return db.session.get(User, int(user_id))

def openpty_with_fallback():
    """Open a PTY with fallback to different device paths for systems where /dev/pty doesn't exist"""
    # First try the standard pty.openpty()
    try:
        master, slave = pty.openpty()
        return master, slave
    except OSError as e:
        if debug:
            print(f"[DEBUG] Standard pty.openpty() failed: {e}, trying fallback methods")

    # Fallback: try to open /dev/ptmx directly
    ptmx_paths = ['/dev/ptmx', '/dev/pts/ptmx']
    for ptmx_path in ptmx_paths:
        try:
            if os.path.exists(ptmx_path):
                # Open master PTY
                master = os.open(ptmx_path, os.O_RDWR | os.O_NOCTTY)
                if master < 0:
                    continue

                # Get slave PTY name
                slave_name = os.ttyname(master)
                if not slave_name:
                    os.close(master)
                    continue

                # Open slave PTY
                slave = os.open(slave_name, os.O_RDWR | os.O_NOCTTY)
                if slave < 0:
                    os.close(master)
                    continue

                if debug:
                    print(f"[DEBUG] Successfully opened PTY using {ptmx_path}: master={master}, slave={slave}")
                return master, slave
        except (OSError, AttributeError) as e:
            if debug:
                print(f"[DEBUG] Failed to open PTY using {ptmx_path}: {e}")
            continue

    # Last resort: try to create PTY devices manually
    try:
        # Try to find an available PTY number
        for i in range(256):  # Try PTY numbers 0-255
            pty_name = f"/dev/pts/{i}"
            try:
                if os.path.exists(pty_name):
                    continue
                # Try to create the PTY device
                master = os.open('/dev/ptmx', os.O_RDWR | os.O_NOCTTY)
                slave_name = os.ttyname(master)
                if slave_name and os.path.exists(slave_name):
                    slave = os.open(slave_name, os.O_RDWR | os.O_NOCTTY)
                    if debug:
                        print(f"[DEBUG] Created PTY manually: master={master}, slave={slave}")
                    return master, slave
                os.close(master)
            except (OSError, AttributeError):
                continue
    except Exception as e:
        if debug:
            print(f"[DEBUG] Manual PTY creation failed: {e}")

    # If all methods fail, raise the original exception
    raise OSError("Failed to open PTY: no available PTY devices found")

# Create database and default admin user
with app.app_context():
    db.create_all()
    if not User.query.filter_by(username='admin').first():
        admin = User(username='admin', password_hash=generate_password_hash('admin123'), is_admin=True)
        db.session.add(admin)
        db.session.commit()

# Flask routes
@app.route('/')
@login_required
def index():
    global args
    return render_template('index.html',
                         clients=list(clients.keys()),
                         websocket_port=args.port,
                         domain=args.domain)

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        user = User.query.filter_by(username=username).first()
        if user and check_password_hash(user.password_hash, password):
            login_user(user)
            return redirect(url_for('index'))
        flash('Invalid username or password')
    return render_template('login.html')

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

@app.route('/users')
@login_required
def users():
    if not current_user.is_admin:
        flash('Access denied')
        return redirect(url_for('index'))
    users = User.query.all()
    return render_template('users.html', users=users)

@app.route('/add_user', methods=['POST'])
@login_required
def add_user():
    if not current_user.is_admin:
        return jsonify({'error': 'Access denied'}), 403

    username = request.form.get('username')
    password = request.form.get('password')
    is_admin = request.form.get('is_admin') == 'on'

    if User.query.filter_by(username=username).first():
        return jsonify({'error': 'Username already exists'}), 400

    user = User(username=username, password_hash=generate_password_hash(password), is_admin=is_admin)
    db.session.add(user)
    db.session.commit()
    return jsonify({'success': True})

@app.route('/edit_user/<int:user_id>', methods=['POST'])
@login_required
def edit_user(user_id):
    if not current_user.is_admin:
        return jsonify({'error': 'Access denied'}), 403

    user = User.query.get_or_404(user_id)
    user.username = request.form.get('username')
    user.is_admin = request.form.get('is_admin') == 'on'

    password = request.form.get('password')
    if password:
        user.password_hash = generate_password_hash(password)

    db.session.commit()
    return jsonify({'success': True})

@app.route('/delete_user/<int:user_id>', methods=['POST'])
@login_required
def delete_user(user_id):
    if not current_user.is_admin:
        return jsonify({'error': 'Access denied'}), 403

    user = User.query.get_or_404(user_id)
    db.session.delete(user)
    db.session.commit()
    return jsonify({'success': True})

@app.route('/terminal/<client_id>')
@login_required
def terminal(client_id):
    if client_id not in clients:
        flash('Client not connected')
        return redirect(url_for('index'))
    return render_template('terminal.html', client_id=client_id)

@app.route('/api/clients')
@login_required
def get_clients():
    return jsonify({
        'clients': list(clients.keys()),
        'count': len(clients)
    })

@app.route('/logos/<path:filename>')
def logos_files(filename):
    return send_from_directory('logos', filename)

@app.route('/terminal/<client_id>/connect', methods=['POST'])
@login_required
def connect_terminal(client_id):
    username = request.form.get('username', 'root')
    request_id = str(uuid.uuid4())
    command = ['wsssh', '-p', str(args.port), f'{username}@{client_id}.{args.domain}']
    # Spawn wsssh process with pty using fallback method
    master, slave = openpty_with_fallback()
    slave_name = os.ttyname(slave)

    def set_controlling_terminal():
        os.setsid()

        # Set the controlling terminal
        try:
            fcntl.ioctl(slave, termios.TIOCSCTTY, 0)
        except (OSError, AttributeError):
            pass  # Some systems don't support TIOCSCTTY

        # Set terminal size to match xterm.js dimensions
        import struct
        winsize = struct.pack('HHHH', 24, 80, 0, 0)
        try:
            fcntl.ioctl(0, termios.TIOCSWINSZ, winsize)
        except (OSError, AttributeError):
            pass

        # Set raw mode - let SSH client handle terminal behavior
        import tty
        try:
            tty.setraw(0)
        except (OSError, AttributeError):
            pass

    proc = subprocess.Popen(
        command,
        stdin=slave,
        stdout=slave,
        stderr=slave,
        preexec_fn=set_controlling_terminal,
        env=dict(os.environ, TERM='xterm', COLUMNS='80', LINES='24')
    )
    os.close(slave)

    # Start a thread to read output
    output_buffer = []
    def read_output():
        output_buffer.append(f'Process PID: {proc.pid}\r\n')
        while proc.poll() is None:
            r, w, e = select.select([master], [], [], 0.1)
            if master in r:
                try:
                    data = os.read(master, 1024)
                    if data:
                        decoded = data.decode('utf-8', errors='ignore')
                        output_buffer.append(decoded)
                except:
                    break
        # Read any remaining data
        try:
            data = os.read(master, 1024)
            while data:
                decoded = data.decode('utf-8', errors='ignore')
                output_buffer.append(decoded)
                data = os.read(master, 1024)
        except:
            pass
        output_buffer.append('\r\nProcess finished.\r\n')
        os.close(master)

    thread = threading.Thread(target=read_output, daemon=True)
    thread.start()

    # Process status will be checked in the reading thread
    active_terminals[request_id] = {'client_id': client_id, 'username': username, 'proc': proc, 'output_buffer': output_buffer, 'master': master}
    return jsonify({'request_id': request_id, 'command': ' '.join(command)})


@app.route('/terminal/<client_id>/data', methods=['GET', 'POST'])
@login_required
def terminal_data(client_id):
    if request.method == 'POST':
        request_id = request.form.get('request_id')
        data = request.form.get('data')
        if request_id in active_terminals:
            proc = active_terminals[request_id]['proc']
            master = active_terminals[request_id]['master']
            if proc.poll() is None:  # Process is still running
                try:
                    os.write(master, data.encode())
                except:
                    pass
        return 'OK'
    else:
        request_id = request.args.get('request_id')
        if request_id in active_terminals:
            proc = active_terminals[request_id]['proc']
            output_buffer = active_terminals[request_id]['output_buffer']
            if output_buffer:
                data = ''.join(output_buffer)
                output_buffer.clear()
                return data
            elif proc.poll() is not None:
                # Process terminated
                return '\r\nProcess terminated.\r\n'
        return ''


@app.route('/terminal/<client_id>/disconnect', methods=['POST'])
@login_required
def disconnect_terminal(client_id):
    request_id = request.form.get('request_id')
    if request_id in active_terminals:
        proc = active_terminals[request_id]['proc']
        if proc.poll() is None:
            proc.terminate()
            try:
                proc.wait(timeout=5)
            except:
                proc.kill()
        del active_terminals[request_id]
    return 'OK'


async def handle_websocket(websocket, path=None):
    try:
        async for message in websocket:
            if debug: print(f"[DEBUG] WebSocket message received: {message[:100]}...")
            data = json.loads(message)
            if data.get('type') == 'register':
                client_id = data['id']
                client_password = data.get('password', '')
                if client_password == server_password:
                    clients[client_id] = websocket
                    print(f"Client {client_id} registered")
                    await websocket.send(json.dumps({"type": "registered", "id": client_id}))
                else:
                    print(f"Client {client_id} registration failed: invalid password")
                    await websocket.send(json.dumps({"type": "registration_error", "error": "Invalid password"}))
            elif data.get('type') == 'tunnel_request':
                client_id = data['client_id']
                request_id = data['request_id']
                if client_id in clients:
                    # Store tunnel mapping
                    active_tunnels[request_id] = {
                        'client_ws': clients[client_id],
                        'wsssh_ws': websocket,
                        'client_id': client_id
                    }
                    # Forward tunnel request to client
                    await clients[client_id].send(json.dumps({
                        "type": "tunnel_request",
                        "request_id": request_id
                    }))
                    await websocket.send(json.dumps({
                        "type": "tunnel_ack",
                        "request_id": request_id
                    }))
                else:
                    await websocket.send(json.dumps({
                        "type": "tunnel_error",
                        "request_id": request_id,
                        "error": "Client not registered"
                    }))
            elif data.get('type') == 'tunnel_data':
                # Forward tunnel data using active tunnel mapping
                request_id = data['request_id']
                if request_id in active_tunnels:
                    tunnel = active_tunnels[request_id]
                    # Forward to client
                    await tunnel['client_ws'].send(json.dumps({
                        "type": "tunnel_data",
                        "request_id": request_id,
                        "data": data['data']
                    }))
            elif data.get('type') == 'tunnel_response':
                # Forward tunnel response from client to wsssh
                request_id = data['request_id']
                if request_id in active_tunnels:
                    tunnel = active_tunnels[request_id]
                    await tunnel['wsssh_ws'].send(json.dumps({
                        "type": "tunnel_data",
                        "request_id": request_id,
                        "data": data['data']
                    }))
            elif data.get('type') == 'tunnel_close':
                request_id = data['request_id']
                if request_id in active_tunnels:
                    tunnel = active_tunnels[request_id]
                    # Forward close to client
                    await tunnel['client_ws'].send(json.dumps({
                        "type": "tunnel_close",
                        "request_id": request_id
                    }))
                    # Clean up tunnel
                    del active_tunnels[request_id]
    except websockets.exceptions.ConnectionClosed:
        # Remove from registry and clean up tunnels
        disconnected_client = None
        for cid, ws in clients.items():
            if ws == websocket:
                disconnected_client = cid
                del clients[cid]
                print(f"Client {cid} disconnected")
                break

        # Clean up active tunnels for this client
        if disconnected_client:
            tunnels_to_remove = []
            for request_id, tunnel in active_tunnels.items():
                if tunnel['client_id'] == disconnected_client:
                    tunnels_to_remove.append(request_id)
            for request_id in tunnels_to_remove:
                del active_tunnels[request_id]

async def main():
    parser = argparse.ArgumentParser(description='WebSocket SSH Daemon (wssshd)')
    parser.add_argument('--host', required=True, help='WebSocket server host')
    parser.add_argument('--port', type=int, required=True, help='WebSocket server port')
    parser.add_argument('--domain', required=True, help='Base domain name')
    parser.add_argument('--password', required=True, help='Registration password')
    parser.add_argument('--web-host', help='Web interface host (optional)')
    parser.add_argument('--web-port', type=int, help='Web interface port (optional)')
    parser.add_argument('--web-https', action='store_true', help='Enable HTTPS for web interface')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')

    global args
    args = parser.parse_args()
    global debug
    debug = args.debug
    global server_password
    server_password = args.password

    # Load certificate
    if getattr(sys, 'frozen', False):
        # Running as bundled executable
        bundle_dir = sys._MEIPASS
        cert_path = os.path.join(bundle_dir, 'cert.pem')
        key_path = os.path.join(bundle_dir, 'key.pem')
    else:
        # Running as script
        cert_path = 'cert.pem'
        key_path = 'key.pem'

    ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    ssl_context.load_cert_chain(cert_path, key_path)

    # Start WebSocket server
    ws_server = await websockets.serve(handle_websocket, args.host, args.port, ssl=ssl_context)

    print(f"WebSocket SSH Daemon running on {args.host}:{args.port}")

    # Start web interface if specified
    if args.web_host and args.web_port:
        # Handle HTTPS setup
        ssl_context = None
        protocol = "http"

        if args.web_https:
            protocol = "https"
            # Use config directory for web certificates
            web_cert_path = os.path.join(config_dir, 'web-cert.pem')
            web_key_path = os.path.join(config_dir, 'web-key.pem')

            # Generate self-signed certificate if it doesn't exist
            if not os.path.exists(web_cert_path) or not os.path.exists(web_key_path):
                print("Generating self-signed certificate for web interface...")
                os.system(f'openssl req -x509 -newkey rsa:4096 -keyout {web_key_path} -out {web_cert_path} -days 365 -nodes -subj "/C=US/ST=State/L=City/O=Organization/CN=localhost"')

            ssl_context = (web_cert_path, web_key_path)

        def run_flask():
            app.run(host=args.web_host, port=args.web_port, debug=debug, use_reloader=False, threaded=True)

        flask_thread = threading.Thread(target=run_flask, daemon=True)
        flask_thread.start()
        print(f"Web interface available at {protocol}://{args.web_host}:{args.web_port}")

    await ws_server.wait_closed()

if __name__ == '__main__':
    asyncio.run(main())