Add new relic based container exec (#27)

Author: wagoodman

.gitignore

@@ -5,6 +5,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
+.pytest_cache
 
 # C extensions
 *.so

Makefile

@@ -20,7 +20,7 @@ _upload: test
 
 tag:
 	git tag v`cat bridgy/version.py | grep __version__ | awk -F"=" '{print $$2}' | sed -e 's/^[ \t]*//' | tr -d "'"` && \
-	git push --tags 
+	git push --tags
 
 bootstrap3: venv3
 	. venv3/bin/activate

bridgy/__main__.py

@@ -9,6 +9,8 @@
   bridgy init
   bridgy ssh (-t | --tmux) [-adsuvw] [-l LAYOUT] <host>...
   bridgy ssh [-duv] <host>
+  bridgy exec (-t | --tmux) [-adsuvw] [-l LAYOUT] <container>...
+  bridgy exec [-duv] <container>
   bridgy list-inventory
   bridgy list-mounts
   bridgy mount [-duv] <host>:<remotedir>
@@ -21,6 +23,7 @@
 Sub-commands:
   init          create the ~/.bridgy/config.yml
   ssh           ssh into the selected host(s)
+  exec          exec into the selected container(s) (interactive + tty)
   mount         use sshfs to mount a remote directory to an empty local directory
   unmount       unmount one or more host sshfs mounts
   list-mounts   show all sshfs mounts
@@ -41,13 +44,15 @@
 
 Configuration Options are in ~/.bridgy/config.yml
 """
-import sys  
+import sys
 if sys.version_info < (3, 0):
     reload(sys)
     sys.setdefaultencoding('utf8')
 
-import os
 import logging
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+import os
 import inquirer
 from inquirer.themes import Theme
 import coloredlogs
@@ -57,6 +62,7 @@
 
 from bridgy.version import __version__
 from bridgy.command import Ssh, Sshfs, RunAnsiblePlaybook
+from bridgy.inventory import InstanceType
 import bridgy.inventory as inventory
 import bridgy.config as cfg
 import bridgy.tmux as tmux
@@ -67,7 +73,7 @@
 
 
 class CustomTheme(Theme):
-    
+
     def __init__(self):
         super(CustomTheme, self).__init__()
 
@@ -82,19 +88,19 @@ def __init__(self):
         self.Checkbox.selected_icon = '◉ ' #✔⬢◉
         self.Checkbox.selected_color = selected_color
         self.Checkbox.unselected_color = utils.term.normal
-        self.Checkbox.unselected_icon = '○ ' #▢ ○ ⬡ 🞅 ⭘ 🔿 🔾 
+        self.Checkbox.unselected_icon = '○ ' #▢ ○ ⬡ 🞅 ⭘ 🔿 🔾
         self.List.selection_color = selection_color
         self.List.selection_cursor = '‣' # ❯
         self.List.unselected_color = utils.term.normal
 
 THEMER = CustomTheme()
 
-def prompt_targets(question, targets=None, instances=None, multiple=True, config=None):
+def prompt_targets(question, targets=None, instances=None, multiple=True, config=None, type=InstanceType.ALL):
     if targets == None and instances == None or targets != None and instances != None:
         raise RuntimeError("Provide exactly one of either 'targets' or 'instances'")
 
     if targets:
-        instances = inventory.search(config, targets)
+        instances = inventory.search(config, targets, type=type)
 
     if len(instances) == 0:
         return []
@@ -143,17 +149,56 @@ def prompt_targets(question, targets=None, instances=None, multiple=True, config
     return selected_hosts
 
 
+@utils.SupportedPlatforms('linux', 'windows', 'osx')
+def exec_handler(args, config):
+    if config.dig('inventory', 'update_at_start') or args['-u']:
+        update_handler(args, config)
+
+    if args ['--tmux'] or config.dig('ssh', 'tmux'):
+        question = "What containers would you like to exec into?"
+        targets = prompt_targets(question, targets=args['<container>'], config=config, type=InstanceType.CONTAINER)
+    else:
+        question = "What containers would you like to exec into?"
+        targets = prompt_targets(question, targets=args['<container>'], config=config, type=InstanceType.CONTAINER, multiple=False)
+
+    if len(targets) == 0:
+        logger.info("No matching instances found")
+        sys.exit(1)
+
+    for instance in targets:
+        if instance.container_id == None:
+            logger.info("Could not find container id for instance: %s" % instance)
+            sys.exit(1)
+
+    commands = collections.OrderedDict()
+    for idx, instance in enumerate(targets):
+        name = '{}-{}'.format(instance.name, idx)
+        commands[name] = Ssh(config, instance, command="sudo -i docker exec -ti %s bash" % instance.container_id).command
+
+    layout = None
+    if args['--layout']:
+        layout = args['--layout']
+
+    if args['--tmux'] or config.dig('ssh', 'tmux'):
+        tmux.run(config, commands, args['-w'], layout, args['-d'], args['-s'])
+    else:
+        cmd = list(commands.values())[0]
+        if args['-d']:
+            logger.debug(cmd)
+        else:
+            os.system(cmd)
+
 @utils.SupportedPlatforms('linux', 'windows', 'osx')
 def ssh_handler(args, config):
     if config.dig('inventory', 'update_at_start') or args['-u']:
         update_handler(args, config)
 
     if args ['--tmux'] or config.dig('ssh', 'tmux'):
         question = "What instances would you like to ssh into?"
-        targets = prompt_targets(question, targets=args['<host>'], config=config)
+        targets = prompt_targets(question, targets=args['<host>'], config=config, type=InstanceType.VM)
     else:
         question = "What instance would you like to ssh into?"
-        targets = prompt_targets(question, targets=args['<host>'], config=config, multiple=False)
+        targets = prompt_targets(question, targets=args['<host>'], config=config, type=InstanceType.VM, multiple=False)
 
     if len(targets) == 0:
         logger.info("No matching instances found")
@@ -257,12 +302,12 @@ def unmount_handler(args, config):
 @utils.SupportedPlatforms('linux', 'windows', 'osx')
 def list_inventory_handler(args, config):
     instances = []
-    for ip, name, aliases, source in inventory.instances(config):
-        if aliases:
-            instances.append( (ip, name, '\n'.join(aliases), source) )
+    for instance in sorted(inventory.instances(config)):
+        if instance.aliases:
+            instances.append( (instance.name, instance.address, '\n'.join(instance.aliases), instance.source, instance.type) )
         else:
-            instances.append( (ip, name, '--- None ---', source) )
-    logger.info(tabulate(instances, headers=['Name', 'Address/Dns', 'Aliases', 'Source']))
+            instances.append( (instance.name, instance.address, '--- None ---', instance.source, instance.type) )
+    logger.info(tabulate(instances, headers=['Name', 'Address/Dns', 'Aliases', 'Source', 'Type']))
 
 
 @utils.SupportedPlatforms('linux', 'windows', 'osx')
@@ -328,6 +373,7 @@ def main():
 
     opts = {
         'ssh': ssh_handler,
+        'exec': exec_handler,
         'mount': mount_handler,
         'list-mounts': list_mounts_handler,
         'list-inventory': list_inventory_handler,

bridgy/command/ssh.py

@@ -3,14 +3,15 @@
 
 class Ssh(object):
 
-    def __init__(self, config, instance):
+    def __init__(self, config, instance, command=''):
         if not hasattr(config, '__getitem__'):
             raise BadConfigError
         if not isinstance(instance, tuple):
             raise BadInstanceError
 
         self.config = config
         self.instance = instance
+        self.custom_command = command
 
     @property
     def destination(self):
@@ -34,11 +35,12 @@ def options(self):
 
         options = self.config.dig('ssh', 'options') or ''
 
-        return '{} {}'.format(bastion, options)
+        return '{} {} -t'.format(bastion, options)
 
 
     @property
     def command(self):
-        cmd = 'ssh {options} {destination}'
+        cmd = 'ssh {options} {destination} {command}'
         return cmd.format(destination=self.destination,
-                          options=self.options )
+                          options=self.options,
+                          command=self.custom_command)

bridgy/inventory/__init__.py

@@ -6,7 +6,7 @@
 
 from bridgy.utils import memoize
 from bridgy.error import MissingBastionHost
-from bridgy.inventory.source import Bastion, Instance, InventorySet
+from bridgy.inventory.source import Bastion, Instance, InventorySet, InstanceType
 from bridgy.inventory.aws import AwsInventory
 from bridgy.inventory.flatfile import CsvInventory
 from bridgy.inventory.newrelic import NewRelicInventory
@@ -28,13 +28,13 @@ def inventory(config):
 
     for source, srcCfg in config.sources():
         if source == 'aws':
-            # the cache directory for the original v1 config did not separate 
+            # the cache directory for the original v1 config did not separate
             # out multiple aws profiles into subdirectories
             if config.version == 1:
                 cache_dir = config.inventoryDir(AwsInventory.name)
             else:
                 cache_dir = config.inventoryDir(AwsInventory.name, srcCfg['name'])
-            
+
             if not os.path.exists(cache_dir):
                 os.mkdir(cache_dir)
 
@@ -71,9 +71,9 @@ def inventory(config):
             inv = NewRelicInventory(data_path=config.inventoryDir(NewRelicInventory.name),
                                     proxies=proxies,
                                     **srcCfg)
- 
+
             inventorySet.add(inv)
-        
+
     return inventorySet
 
 def instance_filter(instance, include_re=None, exclude_re=None):
@@ -110,7 +110,7 @@ def instances(config):
 @memoize
 def get_bastion(config, instance):
     bastion = None
-    
+
     for inv in inventory(config).inventories:
         if inv.name == instance.source and inv.bastion != None:
             bastion = inv.bastion
@@ -132,7 +132,7 @@ def get_bastion(config, instance):
 
     return bastion
 
-def search(config, targets):
+def search(config, targets, type=InstanceType.ALL):
     fuzzy = False
     if config.dig('inventory', 'fuzzy_search'):
         fuzzy = config.dig('inventory', 'fuzzy_search')
@@ -145,7 +145,11 @@ def search(config, targets):
 
     matched_instances = inventory(config).search(targets, fuzzy=fuzzy)
     config_instance_filter = partial(instance_filter, include_re=include_re, exclude_re=exclude_re)
-    return list(filter(config_instance_filter, matched_instances))
+    filtered_instances = list(filter(config_instance_filter, matched_instances))
+    if type == InstanceType.ALL:
+        return filtered_instances
+    else:
+        return [x for x in filtered_instances if x.type == type]
 
 def update(config):
     inventory(config).update()

bridgy/inventory/aws.py

@@ -3,7 +3,7 @@
 import placebo
 import logging
 
-from bridgy.inventory.source import InventorySource, Instance
+from bridgy.inventory.source import InventorySource, Instance, InstanceType
 
 logger = logging.getLogger()
 
@@ -82,9 +82,9 @@ def instances(self):
                 # take note of this instance
                 if name != None and address != None:
                     if len(aliases) > 0:
-                        instances.append(Instance(name, address, tuple(aliases), self.name))
+                        instances.append(Instance(name, address, tuple(aliases), self.name, None, InstanceType.VM))
                     else:
-                        instances.append(Instance(name, address, self.name))
+                        instances.append(Instance(name, address, None, self.name, None, InstanceType.VM))
 
         return instances
 

bridgy/inventory/flatfile.py

@@ -3,7 +3,7 @@
 import sys
 import logging
 
-from bridgy.inventory.source import InventorySource, Instance
+from bridgy.inventory.source import InventorySource, Instance, InstanceType
 
 logger = logging.getLogger()
 
@@ -29,7 +29,7 @@ def instances(self):
             with open(self.csv_path, 'r') as csv_file:
                 reader = csv.DictReader(csv_file, fieldnames=self.fields, delimiter=self.delimiter)
                 for row in reader:
-                    instances.add(Instance(row['name'].strip(), row['address'].strip(), None, self.name))
+                    instances.add(Instance(row['name'].strip(), row['address'].strip(), None, self.name, None, InstanceType.VM))
         except IOError as ex:
             logger.error("Unable to read inventory: %s" % ex)
             sys.exit(1)

bridgy/inventory/newrelic.py

@@ -6,7 +6,8 @@
 except ImportError:
     from urllib import quote_plus
 
-from bridgy.inventory.source import InventorySource, Instance
+from bridgy.inventory.source import InventorySource, Instance, InstanceType
+from bridgy.utils import parseIpFromHostname
 
 class NewRelicInventory(InventorySource):
 
@@ -22,7 +23,8 @@ def __init__(self, account_number, insights_query_api_key, data_path, proxies=No
         self.account_number = account_number
         self.insights_query_api_key = insights_query_api_key
         self.data_file = os.path.join(data_path, '%s.json' % str(account_number))
-        self.query = quote_plus("SELECT entityName, fullHostname, hostname, ipV4Address from NetworkSample LIMIT 999")
+        self.queryVms = quote_plus("SELECT entityName, fullHostname, hostname, ipV4Address from NetworkSample LIMIT 999")
+        self.queryContainers = quote_plus("SELECT containerName, containerId, hostname FROM ProcessSample WHERE containerName IS NOT NULL LIMIT 999")
         if proxies:
             self.proxies = proxies
         else:
@@ -32,24 +34,42 @@ def update(self):
         headers = {'X-Query-Key': self.insights_query_api_key,
                    'Accept': 'application/json'}
 
-        response = requests.get(NewRelicInventory.url.format(self.account_number, self.query),
+        responseVms = requests.get(NewRelicInventory.url.format(self.account_number, self.queryVms),
+                                headers=headers,
+                                proxies=self.proxies)
+
+        responseContainers = requests.get(NewRelicInventory.url.format(self.account_number, self.queryContainers),
                                 headers=headers,
                                 proxies=self.proxies)
 
         with open(self.data_file, 'w') as data_file:
-            data_file.write(response.text)
+            data_file.write(
+                json.dumps({
+                    InstanceType.VM: json.loads(responseVms.text),
+                    InstanceType.CONTAINER: json.loads(responseContainers.text)
+                })
+            )
 
     def instances(self):
         instances = set()
         with open(self.data_file, 'r') as data_file:
             data = json.load(data_file)
 
-        for results_dict in data['results']:
+        for results_dict in data[InstanceType.VM]['results']:
             for event_dict in results_dict['events']:
                 hostname = event_dict['hostname']
                 address = event_dict['ipV4Address'].strip().split("/")[0]
-                if hostname == None:
+                if hostname is None:
                     hostname = address
-                instances.add(Instance(hostname, address, None, self.name))
+                instances.add(Instance(hostname, address, None, self.name, None, InstanceType.VM))
+
+        for results_dict in data[InstanceType.CONTAINER]['results']:
+            for event_dict in results_dict['events']:
+                container_name = event_dict['containerName']
+                container_id = event_dict['containerId']
+                hostname = event_dict['hostname']
+                address = parseIpFromHostname(hostname)
+
+                instances.add(Instance(container_name, address, None, self.name, container_id, InstanceType.CONTAINER))
 
         return list(instances)