webface-server/webface/plugin.ts

import {
  Plugin,
  Configurable,
  EventListener,
  DependencyLoad,
  DependencyUnload
} from '@squeebot/core/lib/plugin';

import express, { RequestHandler } from 'express';
import http from 'http';
import crypto from 'crypto';
import bcrypt from 'bcrypt';
import * as WebSocket from 'ws';

import { logger } from '@squeebot/core/lib/core';
import { EventEmitter } from 'events';
import { Socket } from 'net';

interface RegisteredUser {
  username: string;
  password: string;
}

interface UserRequest {
  command: string;
  arguments?: string[];
  state?: string;
}

interface UserResponse {
  status: 'OK' | 'ERROR';
  list?: any[];
  data?: any;
  command: string;
  state?: string;
  time?: number;
}

class WebSocketServer extends EventEmitter {
  public app = express();
  public server = http.createServer(this.app);
  public wss = new WebSocket.Server({ noServer: true });
  public running = false;

  constructor(
    public port: number,
    public host: string,
    trustProxy = false,
  ) {
    super();
    this.app.disable('x-powered-by');
    if (trustProxy) {
      this.app.set('trust proxy', 1);
    }
  }

  public init(): void {
    this.server.listen(this.port, this.host, () => {
      logger.log('[webface] WebSocket server listening on %s:%s', this.host, this.port);
      this.running = true;
    });
  }

  public destroy(): void {
    this.running = false;
    this.wss.close();
    this.server.close();
  }
}

class WebfaceServer extends WebSocketServer {
  private controlPlugin: any;
  private connections = new Map<string, WebSocket>();
  private issuedTokens = new Map<string, string>();

  constructor(
    public port: number,
    public host: string,
    public authorizedURL: string,
    private users: RegisteredUser[],
    trustProxy = false,
  ) {
    super(port, host, trustProxy);
    this.initializeWebSocket();
    this.initializeAPI();
  }

  /**
   * Set the loaded control plugin and announce that it has been loaded
   * @param control Control plugin
   */
  public setControlPlugin(control: any): void {
    this.controlPlugin = control;
    this.connections.forEach((ws) => {
      ws.send(JSON.stringify({
        status: 'OK_CONTROL',
        message: 'Control plugin has loaded again!'
      }));
    });
  }

  /**
   * Notify websocket connections that the control plugin disappeared
   */
  public controlPluginUnloaded(): void {
    this.controlPlugin = null;
    this.connections.forEach((ws) => {
      ws.send(JSON.stringify({
        status: 'MISSING_CONTROL',
        message: 'Control plugin has unloaded!'
      }));
    });
  }

  /**
   * Get the username of a token or throw an error if missing or invalid
   * @param req HTTP connection
   * @returns Token owner
   */
  public async authorizeAccessToken(req: http.IncomingMessage): Promise<string> {
    const authHead = req.headers.authorization;
    if (!authHead || !authHead.startsWith('Bearer')) {
      throw new Error('Authorization header missing');
    }

    const owner = this.getTokenOwner(authHead.substring(7));

    if (!owner) {
      throw new Error('Unauthorized');
    }

    return owner;
  }

  /**
   * Shut down
   */
  public destroy(): void {
    this.connections.forEach((ws, user) => {
      ws.close();
    });

    this.connections.clear();
    this.issuedTokens.clear();
    super.destroy();
  }

  /**
   * Authorize a token
   * @returns Middleware
   */
  public tokenOwnerMiddleware(): RequestHandler {
    return (req, res, next) => {
      const authHeader = req.get('authorization');
      if (!authHeader || !authHeader.startsWith('Bearer')) {
        return res.status(401).json({
          error: 'unauthorized',
          error_message: 'Unauthorized',
        });
      }

      const token = authHeader.substring(7);
      const tokenOwner = this.getTokenOwner(token);
      if (!tokenOwner) {
        return res.status(401).json({
          error: 'unauthorized',
          error_message: 'Unauthorized',
        });
      }

      res.locals.token = token;
      res.locals.user = tokenOwner;
      next();
    };
  }

  /**
   * Get a token's owner or undefined if invalid or missing
   * @param token Token from headers
   * @returns Token owner or undefined
   */
  private getTokenOwner(token: string): string | undefined {
    const representative = [...this.issuedTokens.entries()]
      .filter(({ 1: v }) => v === token)
      .map(([k]) => k);

    if (!representative?.length) {
      return;
    }

    return representative[0];
  }

  /**
   * Send a status message with available commands
   * @param ws WebSocket client
   * @param username Client's username
   */
  private sendStatus(ws: WebSocket, username: string): void {
    ws.send(JSON.stringify({
      status: this.controlPlugin ? 'OK' : 'MISSING_CONTROL',
      commands: this.controlPlugin?.listControlCommands() || [],
      user: {
        username,
      }
    }));
  }

  /**
   * Handle input from a WebSocket connection
   * @param raw WebSocket message
   * @param ws WebSocket client
   * @param user Client username
   */
  private handleWSLine(raw: WebSocket.RawData, ws: WebSocket, user: string): void {
    const line = raw.toString('utf8');
    let jCmd: UserRequest;
    try {
      jCmd = JSON.parse(line);
    } catch (e: any) {
      ws.send(JSON.stringify({
        status: 'ERROR', message: e.message
      }));
      return;
    }

    const { command, state } = jCmd;
    if (!command) {
      ws.send(JSON.stringify({
        status: 'ERROR',
        message: 'Unknown or missing command',
        command: '',
        state: jCmd.state,
      }));
      return;
    }

    if (command === 'quit' || command === 'exit') {
      ws.close();
      return;
    }

    if (command === 'help' || command === 'status') {
      this.sendStatus(ws, user);
      return;
    }

    const timeStart = Date.now();
    this.controlPlugin.executeControlCommand(command, jCmd.arguments || []).then(
      (data: any) => {
        const response: UserResponse = {
          status: 'OK',
          command,
          state,
          time: Date.now() - timeStart,
        };

        if (Array.isArray(data)) {
          response.list = data;
        } else {
          response.data = data;
        }
        ws.send(JSON.stringify(response));
      }, (error: Error) => {
        ws.send(JSON.stringify({
          status: 'ERROR',
          message: error.message,
          command,
          state,
          time: Date.now() - timeStart,
        }));
      });
  }

  /**
   * Start listening for websocket connections
   */
  private initializeWebSocket(): void {
    this.wss.on('connection', (ws: WebSocket, request: http.IncomingMessage, user: string) => {
      if (this.connections.has(user)) {
        ws.close();
        return;
      }

      logger.log('[webface] User %s connected via WebSocket', user);

      this.sendStatus(ws, user);
      this.connections.set(user, ws);

      ws.on('message', (data) => this.handleWSLine(data, ws, user));
      ws.on('close', () => {
        if (this.connections.has(user)) {
          this.connections.delete(user);
        }
        logger.log('[webface] User %s WebSocket disconnected', user);
      });
    });

    this.server.on('upgrade', (req, socket, head) => {
      this.authorizeAccessToken(req).then((user) => {
        this.wss.handleUpgrade(req, socket as Socket, head, (ws) => {
          this.wss.emit('connection', ws, req, user);
        });
      }, () => {
        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
        socket.destroy();
      });
    });
  }

  /**
   * Set up the HTTP API
   */
  private initializeAPI(): void {
    this.app.use(express.json() as RequestHandler);
    this.app.use((req, res, next) => {
      res.header('Access-Control-Allow-Origin', this.authorizedURL);
      res.header('Access-Control-Allow-Methods', 'POST, GET, DELETE, OPTIONS');
      res.header('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization');
      res.header('Access-Control-Allow-Credentials', 'true');
      next();
    });

    const router = express.Router();
    const authMiddleware = this.tokenOwnerMiddleware();

    router.get('/', (req, res) => {
      res.json({
        status: 'OK',
        uptime: process.uptime(),
        connections: this.connections.size,
      });
    });

    router.get('/squeebot', authMiddleware, (req, res) => {
      res.json({
        status: this.controlPlugin ? 'OK' : 'MISSING_CONTROL',
        uptime: process.uptime(),
        connections: this.connections.size,
        user: {
          username: res.locals.user,
        },
        commands: this.controlPlugin?.listControlCommands() || [],
      });
    });

    router.post('/authorize', (req, res) => {
      const { body: { username, password } } = req;
      const user = this.users.find(
        (userEntry) => userEntry.username.toLowerCase() === username.toLowerCase()
      );
      if (!user || !username || !password) {
        return res.status(401).json({
          error: 'invalid_username_password',
          error_message: 'Invalid username or password',
        });
      }

      bcrypt.compare(password, user.password).then((success) => {
        if (!success) {
          return res.status(401).json({
            error: 'invalid_username_password',
            error_message: 'Invalid username or password',
          });
        }

        const token = crypto.randomBytes(16).toString('hex');
        this.issuedTokens.set(username, token);
        res.json({
          username, token,
        });
      });
    });

    router.post('/introspect', authMiddleware, (req, res) => {
      res.json({ status: 'OK', message: 'Token valid', user: { username: res.locals.user } });
    });

    router.delete('/token', authMiddleware, (req, res) => {
      const { user } = res.locals;
      this.issuedTokens.delete(user);
      const ws = this.connections.get(user);
      if (ws) {
        ws.close();
      }
      res.json({ status: 'OK', message: 'Logged out' });
    });

    router.post('/:command', authMiddleware, (req, res) => {
      if (!this.controlPlugin) {
        return res.status(500).json({
          error: 'control_plugin_unavailable',
          error_message: 'Control plugin is currently unavailable'
        });
      }

      const command = req.params.command;
      const { body } = req;
      if (body?.arguments && !Array.isArray(body.arguments)) {
        return res.status(400).json({
          error: 'arguments_invalid',
          error_message: 'Arguments have to be passed as an array',
        });
      }

      const timeStart = Date.now();
      this.controlPlugin.executeControlCommand(command, body?.arguments || []).then(
        (data: any) => {
          res.json({
            status: 'OK',
            data,
            time: Date.now() - timeStart,
          });
        }, (error: Error) => {
          res.status(400).json({
            error: `${command}_fail`,
            error_message: error.message,
            time: Date.now() - timeStart,
          });
        });
    });

    this.app.use('/api/v1', router);
  }

  public announcePluginLoaded(plugin: string): void {
    this.connections.forEach((ws) => {
      ws.send(JSON.stringify({
        status: 'PLUGIN_LOADED',
        message: `${plugin} has been loaded`,
        plugin,
      }));
    });
  }

  public announcePluginUnloaded(plugin: string): void {
    this.connections.forEach((ws) => {
      ws.send(JSON.stringify({
        status: 'PLUGIN_UNLOADED',
        message: `${plugin} has been unloaded`,
        plugin,
      }));
    });
  }
}

@Configurable({
  port: 5511,
  host: 'localhost',
  trustedRemote: '*',
  trustProxy: false,
  users: [
    {
      username: 'squeeadmin',
      password: ''
    }
  ]
})
class WebfacePlugin extends Plugin {
  private srv?: WebfaceServer;
  @EventListener('pluginUnload')
  public unloadEventHandler(plugin: string | Plugin): void {
    if (plugin === this.name || plugin === this) {
      this.srv?.destroy();
      this.emit('pluginUnloaded', this);
    }
  }

  @EventListener('pluginLoaded')
  public pluginLoadedEvent(plugin: string | Plugin): void {
    let pluginName = plugin as string;
    if (typeof plugin !== 'string') {
      pluginName = plugin.manifest.name;
    }
    this.srv?.announcePluginLoaded(pluginName);
  }

  @EventListener('pluginUnloaded')
  public pluginUnloadedEvent(plugin: string | Plugin): void {
    let pluginName = plugin as string;
    if (typeof plugin !== 'string') {
      pluginName = plugin.manifest.name;
    }
    this.srv?.announcePluginUnloaded(pluginName);
  }

  @DependencyLoad('control')
  public controlLoaded(control: any): void {
    if (this.srv) {
      this.srv.setControlPlugin(control);
    }
  }

  @DependencyUnload('control')
  public controlUnloaded(): void {
    if (this.srv) {
      this.srv.controlPluginUnloaded();
    }
  }

  initialize(): void {
    const users = this.config.get('users', []) as RegisteredUser[];
    const defaultUser = users.find(
      (user) => user.username === 'squeeadmin'
    );

    // Set default user password
    if (defaultUser && defaultUser.password === '') {
      const defaultPassword = 'squeeadmin13';
      bcrypt.hash(defaultPassword, 10).then((hashed) => {
        defaultUser.password = hashed;
        logger.warn('[webface] !!! The default user is "squeeadmin" and the password has been set to "squeeadmin13" !!!');
        logger.warn('[webface] !!! Change this password ASAP from the front-end of your choice !!!');
        this.config.set('users', users);
        this.config.save().catch((e) => logger.error('[webface] Failed to save users:', e.message));
      });
    }

    // Create the server
    this.srv = new WebfaceServer(
      this.config.get('port') as number,
      this.config.get('host'),
      this.config.get('trustedRemote', '*') as string,
      users,
      this.config.get('trustProxy', false),
    );

    this.srv.init();
  }
}

module.exports = WebfacePlugin;