#!/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
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
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 = {}
debug = False
server_password = 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 User.query.get(int(user_id))

# 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():
    return render_template('index.html', clients=list(clients.keys()))

@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)

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']
                if client_id in clients:
                    # Forward tunnel request to client
                    await clients[client_id].send(json.dumps({
                        "type": "tunnel_request",
                        "request_id": data['request_id']
                    }))
                    await websocket.send(json.dumps({
                        "type": "tunnel_ack",
                        "request_id": data['request_id']
                    }))
                else:
                    await websocket.send(json.dumps({
                        "type": "tunnel_error",
                        "request_id": data['request_id'],
                        "error": "Client not registered"
                    }))
            elif data.get('type') == 'tunnel_data':
                # Forward tunnel data to appropriate client
                client_id = data.get('client_id')
                if client_id and client_id in clients:
                    await clients[client_id].send(json.dumps({
                        "type": "tunnel_data",
                        "request_id": data['request_id'],
                        "data": data['data']
                    }))
            elif data.get('type') == 'tunnel_close':
                client_id = data.get('client_id')
                if client_id and client_id in clients:
                    await clients[client_id].send(json.dumps({
                        "type": "tunnel_close",
                        "request_id": data['request_id']
                    }))
    except websockets.exceptions.ConnectionClosed:
        # Remove from registry
        for cid, ws in clients.items():
            if ws == websocket:
                del clients[cid]
                print(f"Client {cid} disconnected")
                break

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('--debug', action='store_true', help='Enable debug output')

    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:
        def run_flask():
            app.run(host=args.web_host, port=args.web_port, debug=False, use_reloader=False)

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

    await ws_server.wait_closed()

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