Add websocket support for testing rules. Add search endpoint for checking user queries before a test runs.

Author: johnsusek

README.md

@@ -14,7 +14,7 @@ The most convenient way to run the ElastAlert server is by using our Docker cont
 To run the Docker image you will want to mount the volumes for configuration and rule files to keep them after container updates. In order to do that conveniently, please do: `git clone https://github.com/bitsensor/elastalert.git; cd elastalert`
 
 ```bash
-docker run -d -p 3030:3030 \
+docker run -d -p 3030:3030 -p 3333:3333 \
     -v `pwd`/config/elastalert.yaml:/opt/elastalert/config.yaml \
     -v `pwd`/config/config.json:/opt/elastalert-server/config/config.json \
     -v `pwd`/rules:/opt/elastalert/rules \
@@ -60,6 +60,7 @@ You can use the following config options:
 {
   "appName": "elastalert-server", // The name used by the logging framework.
   "port": 3030, // The port to bind to
+  "wsport": 3333, // The port to bind to for websockets
   "elastalertPath": "/opt/elastalert",  // The path to the root ElastAlert folder. It's the folder that contains the `setup.py` script.
   "start": "2014-01-01T00:00:00", // Optional date to start querying from
   "end": "2016-01-01T00:00:00", // Optional date to stop querying at
@@ -208,10 +209,10 @@ This server exposes the following REST API's:
         }
       }
       ``` 
-    
-- **GET `/test_stream`**
 
-    This allows you to test a rule and get a [Server Sent Events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events) event stream back. Pass params `rule` (yaml string) and `options` (JSON string) to start receiving events.
+- **WEBSOCKET `/test`**
+    
+    This allows you to test a rule and receive progress over a websocket. Send a message as JSON object (stringified) with two keys: `rule` (yaml string) and `options` (JSON object). You will receive progress messages over the socket as the test runs.
 
 - **GET `/metadata/:type`**
 

config/config.json

@@ -2,6 +2,7 @@
   "appName": "elastalert-server",
   "port": 3030,
   "elastalertPath": "/opt/elastalert",
+  "wsport": 3333,
   "verbose": false,
   "es_debug": false,
   "debug": false,

package.json

@@ -35,7 +35,8 @@
     "raven": "^2.6.1",
     "request": "^2.85.0",
     "request-promise-native": "^1.0.5",
-    "tar": "^4.4.1"
+    "tar": "^4.4.1",
+    "ws": "^6.0.0"
   },
   "devDependencies": {
     "eslint": "^4.17.0",

src/common/websocket.js

@@ -0,0 +1,27 @@
+import WebSocket from 'ws';
+
+export var wss = null;
+
+export function listen(port) {
+  wss = new WebSocket.Server({ port, path: '/test' });
+
+  wss.on('connection', ws => {
+    ws.isAlive = true;
+    ws.on('pong', () => {
+      ws.isAlive = true;
+    });
+  });
+
+  return wss;
+}
+
+// Keepalive in case clients lose connection during a long rule test.
+// If client doesn't respond in 10s this will close the socket and 
+// therefore stop the elastalert test from continuing to run detached.
+setInterval(() => {
+  wss.clients.forEach(ws => {
+    if (ws.isAlive === false) return ws.terminate();
+    ws.isAlive = false;
+    ws.ping(() => {});
+  });
+}, 10000);

src/controllers/test/index.js

@@ -19,7 +19,7 @@ export default class TestController {
     });
   }
 
-  testRule(rule, options, stream, response) {
+  testRule(rule, options, socket) {
     const self = this;
     let tempFileName = '~' + randomstring.generate() + '.temp';
     let tempFilePath = path.join(self.testFolder, tempFileName);
@@ -55,31 +55,41 @@ export default class TestController {
               break;
           }
 
+
           try {
             let testProcess = spawn('python', processOptions, {
               cwd: self._elastalertPath
             });
 
+            // When the websocket closes we kill the test process
+            // so it doesn't keep running detached
+            if (socket) {
+              socket.on('close', () => {
+                testProcess.kill();
+              });
+            }
+              
             testProcess.stdout.on('data', function (data) {
-              if (stream) {
-                response.write('event: result\ndata: ' + data.toString() + '\n\n');
+              if (socket) {
+                socket.send(JSON.stringify({ 
+                  event: 'result',
+                  data: data.toString() 
+                }));
               }
               stdoutLines.push(data.toString());
             });
 
             testProcess.stderr.on('data', function (data) {
-              if (stream) {
-                response.write('event: progress\ndata: ' + data.toString() + '\n\n');
+              if (socket) {
+                socket.send(JSON.stringify({ 
+                  event: 'progress',
+                  data: data.toString() 
+                }));
               }
               stderrLines.push(data.toString());
             });
 
             testProcess.on('exit', function (statusCode) {
-              if (stream) {
-                response.write('event: done\ndata: DONE\n\n');
-                response.end();
-              }
-
               if (statusCode === 0) {
                 if (options.format === 'json') {
                   resolve(stdoutLines.join(''));
@@ -88,8 +98,10 @@ export default class TestController {
                   resolve(stdoutLines.join('\n'));
                 }
               } else {
-                reject(stderrLines.join('\n'));
-                logger.error(stderrLines.join('\n'));
+                if (!socket) {
+                  reject(stderrLines.join('\n'));
+                  logger.error(stderrLines.join('\n'));  
+                }
               }
 
               fileSystem.deleteFile(tempFilePath)
@@ -106,6 +118,8 @@ export default class TestController {
           logger.error(`Failed to write file ${tempFileName} to ${self.testFolder} with error:`, error);
           reject(error);
         });
+    }).catch((error) => {
+      logger.error('Failed to test rule with error:', error);
     });
   }
 

src/elastalert_server.js

@@ -4,6 +4,7 @@ import Logger from './common/logger';
 import config from './common/config';
 import path from 'path';
 import FileSystem from './common/file_system';
+import { listen } from './common/websocket';
 import setupRouter from './routes/route_setup';
 import ProcessController from './controllers/process';
 import RulesController from './controllers/rules';
@@ -77,8 +78,27 @@ export default class ElastalertServer {
         self._fileSystemController.createDirectoryIfNotExists(self.getDataFolder()).catch(function (error) {
           logger.error('Error creating data folder with error:', error);
         });
-
+        
         logger.info('Server listening on port ' + config.get('port'));
+
+        let wss = listen(config.get('wsport'));
+
+        wss.on('connection', ws => {
+          ws.on('message', (data) => {
+            try {
+              data = JSON.parse(data);
+              if (data.rule) {
+                let rule = data.rule;
+                let options = data.options;
+                self._testController.testRule(rule, options, ws);
+              }
+            } catch (error) {
+              console.log(error);
+            }
+          });
+        });
+
+        logger.info('Websocket listening on port 3333');
       } catch (error) {
         logger.error('Starting server failed with error:', error);
         process.exit(1);

src/handlers/test/stream.js

@@ -12,7 +12,6 @@ import templateGetHandler from '../handlers/templates/id/get';
 import templatePostHandler from '../handlers/templates/id/post';
 import templateDeleteHandler from '../handlers/templates/id/delete';
 import testPostHandler from '../handlers/test/post';
-import testStreamGetHandler from '../handlers/test/stream';
 import configGetHandler from '../handlers/config/get';
 import configPostHandler from '../handlers/config/post';
 import metadataHandler from '../handlers/metadata/get';
@@ -70,11 +69,6 @@ let routes = [
     method: 'POST',
     handler: testPostHandler
   }, 
-  {
-    path: 'test_stream',
-    method: 'GET',
-    handler: testStreamGetHandler
-  },
   {
     path: 'config',
     method: ['GET', 'POST'],