Merge pull request #17 from criteo-forks/configure-env

Add support for custom environment in Mesos term

README.md

@@ -86,6 +86,7 @@ Here are the details of available options.
 | Parameter                               | Description                                                                              |
 |-----------------------------------------|------------------------------------------------------------------------------------------|
 | MESOS\_TERM\_COMMAND                    | The command to be run in the container when a user logs in. (Default: /bin/sh)           |
+| MESOS\_TERM\_ENVIRONMENT                | List of environment variable to enrich the shell with (NAME=value, colon separated)      |
 | MESOS\_TERM\_ENABLE\_PER\_APP\_ADMINS   | If 'true', application administrators can be declared with the Mesos label MESOS\_TERM\_DEBUG\_GRANTED\_TO label. It means those users can log into the application containers. An example is provided below. (Default: false) |
 | MESOS\_TERM\_ALLOWED\_TASK\_ADMINS      | White list of application administrators (users or groups) allowed to override application configuration through Mesos label |
 | MESOS\_TERM\_ENABLE\_RIGHTS\_DELEGATION | If 'true', super administrators can delegate rights to log into one specific container to one person for a certain amount of time. (Default: false) |

scripts/tests/docker-compose.yml

@@ -132,6 +132,7 @@ services:
       - MESOS_TERM_MESOS_STATE_CACHE_TIME=60
       - MESOS_TERM_SESSION_SECRET=abcd
       - MESOS_TERM_SUPER_ADMINS=admins
+      - 'MESOS_TERM_ENVIRONMENT=TOTO=tata:Bad=var'
 
   mesos-term-no-auth:
     build:

src/controllers/terminal.ts

@@ -106,6 +106,11 @@ function spawnTerminal(
     params.push(task.user);
   }
 
+  if (env.EXTRA_ENV) {
+    params.push('--env');
+    params.push(env.EXTRA_ENV);
+  }
+
   const options: NodePty.IPtyForkOptions = {
     name: 'bash'
   };

src/env_vars.ts

@@ -31,6 +31,7 @@ export interface EnvVars {
   ALLOWED_TASK_ADMINS: string[];
   ENABLE_PER_APP_ADMINS?: boolean;
   ENABLE_RIGHTS_DELEGATION?: boolean;
+  EXTRA_ENV: string;
   COMMAND: string;
 }
 
@@ -58,6 +59,7 @@ export const env: EnvVars = {
   MESOS_MASTER_URL: getOrExit('MESOS_TERM_MESOS_MASTER_URL'),
   AUTHORIZATIONS_ENABLED: authorizations_enabled,
   MESOS_STATE_CACHE_TIME: parseFloat(getOrExit('MESOS_TERM_MESOS_STATE_CACHE_TIME')),
+  EXTRA_ENV: getOrElse('MESOS_TERM_ENVIRONMENT', ''),
   COMMAND: getOrElse('MESOS_TERM_COMMAND', '/bin/sh')
 };
 

src/python/terminal.py

@@ -8,6 +8,7 @@
 import time
 import uuid
 import urllib.parse
+import re
 
 from functools import partial
 from queue import Queue
@@ -31,6 +32,8 @@ class TaskIO(object):
     :type cmd: str
     :param args: Additional arguments for the command
     :type args: str
+    :param env: List of environment variable to enrich the shell with (NAME=value, colon separated)
+    :type env: str
     :param interactive: whether to attach STDIN of the current
                         terminal to the new command being launched
     :type interactive: bool
@@ -45,9 +48,11 @@ class TaskIO(object):
     HEARTBEAT_INTERVAL_NANOSECONDS = HEARTBEAT_INTERVAL * 1000000000
 
     def __init__(self, agent_url, parent_container_id, user=None, cmd=None,
+                 env=None,
                  args=None, interactive=False, tty=False):
         # Store relevant parameters of the call for later.
         self.cmd = cmd
+        self.env = env
         self.interactive = interactive
         self.tty = tty
         self.args = args
@@ -228,6 +233,21 @@ def _launch_nested_container_session(self):
                     'value': self.cmd,
                     'arguments': [self.cmd] + self.args,
                     'shell': False}}}
+        if self.env is not None:
+            env_vars = []
+            env_var_regex = re.compile('^([A-Z_][A-Z0-9_]*)=(.*)$')
+            for env_var in self.env.split(':'):
+                matches = env_var_regex.match(env_var)
+                if matches and len(matches.groups()) == 2:
+                    env_vars.append({
+                        'name': matches.group(1),
+                        'type': 'VALUE',
+                        'value': matches.group(2)
+                    })
+            message['launch_nested_container_session']['command']['environment'] = {
+                'variables': env_vars
+            }
+
         if not self.user is None:
           message['launch_nested_container_session']['command']['user'] =\
              self.user
@@ -489,10 +509,12 @@ def _window_resize(self, signum, frame):
   parser.add_argument('container_id', type=str, help='The container id to connect to.')
   parser.add_argument('--user', type=str, help='The user to run the command as.')
   parser.add_argument('--cmd', type=str, default="/bin/sh", help='The command to run in the container.')
+  parser.add_argument('--env', type=str, help='List of environment variable to enrich the shell with (NAME=value, colon separated).')
   args = parser.parse_args()
 
   t = TaskIO(agent_url=args.agent_url, parent_container_id=args.container_id,
              tty=True, user=args.user, interactive=True, cmd=args.cmd,
+             env=args.env,
              args=[])
   t.run()
 

tests/apps_helpers.ts

@@ -26,6 +26,27 @@ export function interactWithTerminal(driver: any) {
     });
 }
 
+export function testAnswerToCommand(driver: any, command: string, expected: string) {
+  return driver.wait(webdriver.until.elementLocated(webdriver.By.css('.xterm-rows>div')), TIMEOUT)
+    .then(function(el: webdriver.WebElement) {
+      return Bluebird.resolve(driver.wait(webdriver.until.elementTextContains(el, 'runs'), TIMEOUT));
+    })
+    .then(function() {
+      return Bluebird.resolve(
+        driver.wait(webdriver.until.elementLocated(webdriver.By.css('.terminal')), TIMEOUT));
+    })
+    .then(function(el: webdriver.WebElement) {
+      return Bluebird.resolve(el.sendKeys(command + '\n'));
+    })
+    .then(function() {
+      return Bluebird.resolve(
+        driver.wait(webdriver.until.elementLocated(webdriver.By.css('.xterm-rows')), TIMEOUT));
+    })
+    .then(function(el: webdriver.WebElement) {
+      return Bluebird.resolve(driver.wait(webdriver.until.elementTextContains(el, expected), TIMEOUT));
+    });
+}
+
 export function checkInteractionsWithTerminal(
   port: number,
   user: string,

tests/term_env.ts

@@ -0,0 +1,37 @@
+import Bluebird = require('bluebird');
+import AppsHelpers = require('./apps_helpers');
+import helpers = require('./helpers');
+import webdriver = require('selenium-webdriver');
+
+
+describe('custom environment variables', function() {
+  const user = 'john';
+  const port = 3000;
+  const appName = 'app1';
+
+  describe('environment variable TOTO=tata is configured', function() {
+    it('environment variable is set', function() {
+      this.timeout(5000);
+      const instanceId = this.mesosTaskIds[appName];
+      return helpers.withChrome(function(driver) {
+        return Bluebird.resolve(driver.get(`http://${user}:password@localhost:${port}/login/${instanceId}`))
+          .then(function() {
+            return AppsHelpers.testAnswerToCommand(driver, 'echo $TOTO', 'tata');
+          });
+      });
+    });
+  });
+
+  describe('environment variable Bad=var is configured', function() {
+    it('environment variable is not set', function() {
+      this.timeout(5000);
+      const instanceId = this.mesosTaskIds[appName];
+      return helpers.withChrome(function(driver) {
+        return Bluebird.resolve(driver.get(`http://${user}:password@localhost:${port}/login/${instanceId}`))
+          .then(function() {
+            return AppsHelpers.testAnswerToCommand(driver, 'echo ${Bad-not set}', 'not set');
+          });
+      });
+    });
+  });
+});