websocket-python.js

const WebSocket = require('ws');
const PythonShell = require('python-shell');
const tmp = require('tmp');
const shortid = require('shortid');
const fs = require('fs');
const mkdirp = require('mkdirp');

// TODO clean up code, error handling, watchdog timer

const wss = new WebSocket.Server({ port: 8080 });

console.log('Started WS Python server on 8080');
wss.on('connection', function connection(ws, req) {
  const wsId = shortid.generate();
  let pyShell = null;

  console.log('Connection from', req.connection.remoteAddress);

  ws.on('message', function incoming(data) {
    const payload = JSON.parse(data);

    if (payload.action === 'run') {
      console.log("Creating dir at ", `/tmp/${wsId}-XXXXXX`);
      tmp.dir({ template: `/tmp/${wsId}-XXXXXX`, unsafeCleanup: true }, (err, path) => {
        console.log("Created", path, err);
        let numFiles = 0;
        let filesSaved = 0;
        const setupShell = () => {
          if (filesSaved === numFiles) {
            console.log(payload);
            console.log("Setting up shell...", path + '/' + payload.entryPoint);
            if (pyShell && !pyShell.terminated) {
              ws.send(JSON.stringify({
                error: 'Shell still running'
              }));
            } else {
              pyShell = new PythonShell(payload.entryPoint, {
                scriptPath: path,
                args: payload.args || []
              });
              console.log('set up shell')
              pyShell.on('message', (message) => {
                ws.send(JSON.stringify({
                  action: 'message',
                  message: message
                }));
              });

              pyShell.on('error', (pyErr) => {
                ws.send(JSON.stringify({
                  action: 'error',
                  message: pyErr.traceback.replace(path, '')
                }));
              });
            }
          }
        };
        
        for (let filePath in payload.files) {
          const fullPath = path + '/' + filePath;
          const pathParts = fullPath.split('/');
          const pathPrefix = pathParts.slice(0, pathParts.length-1).join('/');

          numFiles++;
          mkdirp(pathPrefix, (mkdirErr) => {
            fs.writeFile(fullPath, payload.files[filePath], (writeErr) => {
              console.log("wrote", fullPath, writeErr);
              filesSaved++;
              setupShell();
            });
          });
        }
      });
    } else if (payload.action === 'message') {
      if (!pyShell) {
        ws.send(JSON.stringify({error: 'Shell not initialised'}));
      } else {
        pyShell.send(payload.message);
      }
    } else if (payload.action === 'end') {
      pyShell.end((err) => {
        ws.send(JSON.stringify({
          action: 'end'
        }));
      });
    }
  });
});