logs.py

#!/usr/bin/env python
"""
Access Elasticsearch based logs.
"""

import datetime
import random
import requests
import shellish
import time
from . import s3, base

URL_FORMAT = 'https://logs.infra.%(region)s.homecu.io'


class ES(object):
    """ Wrap API access to elasticsearch. """

    def __init__(self, url):
        self._url = url.rstrip('/')
        self._http = self._establish()

    def _establish(self):
        http = requests.Session()
        # XXX: Keep this code internal please.
        _1 = 443, 444, 457, 440, 447, 458, 438, 448, 451, 442
        _2 = 400, 392, 461, 444, 448, 443, 444, 452
        # Use the redirection as our auth success check.  There is no
        # 4xx level response to auth failures with oauth2_proxy inside
        # via nginx's auth-request.

        def unfurl(arr):
            return (chr(x - (7 ** 3)) for x in reversed(arr))
        success_url = '/__auth_success_%d__' % (random.random() * 2 ** 64)
        r = http.post(self._url + '/oauth2/sign_in', data={
            "username": ''.join(unfurl(_1)),
            "password": ''.join(unfurl(_2)),
            "rd": success_url
        }, allow_redirects=False)
        if r.status_code not in (301, 302) or \
           r.headers.get('location') != success_url:
            raise Exception("Login failure: [%d] %s" % (r.status_code, r.text))
        return http

    def search(self, index, type=None, fields=None, start=None, limit=None,
               sort_order='asc', **bool_filters):
        if type is not None:
            urn = '%s/%s' % (index, type)
        else:
            urn = index
        params = {
            "query": {
                "constant_score": {
                    "filter": {
                        "bool": bool_filters
                    }
                }
            },
            "sort": [{
                "timestamp": {"order": sort_order}
            }, {
                "_uid": {"order": "desc"}  # tiebreaker/seen-tracker
            }]
        }
        if start is not None:
            params['from'] = start
        if limit is not None:
            params['size'] = limit
        if fields is not None:
            params['_source'] = fields
        r = self._http.get('%s/%s/_search?pretty' % (self._url, urn),
                           json=params)
        r.raise_for_status()
        data = r.json()
        assert not data['_shards']['failed']
        assert not data['timed_out']
        return data['hits']['hits']


class LogsCommand(s3.S3BucketMixin, base.HostingCommand):
    """ Dump or follow Odyssey system logs. """

    name = 'logs'
    levels = (
        "debug",
        "info",
        "warning",
        "error",
        "critical",
    )

    def setup_args(self, parser):
        self.add_stack_argument()
        self.table_args_parser = self.add_table_arguments()
        self.add_argument('--follow', '-f', action='store_true',
                          help='Follow mode.')
        self.add_argument('--level', '-l', choices=self.levels,
                          help='Minimum log level.')
        self.add_argument('--image', '-i', help='Only show logs for this '
                          'image.')
        self.add_argument('--highlight', '-hl', nargs='+',
                          help='Highlight phrases found in log messages.')
        self.add_argument('--query', '-q', nargs='+',
                          help='Elasticsearch Query [EXPERT USE]')
        self.add_argument('-n', type=int, default=20,
                          help='Number of lines to initially show.')
        self.add_argument('--verbose', '-v', action='store_true',
                          help='Verbose log output.')
        self.add_argument('--url-override', default=None,
                          help='Override the Elasticsearch URL.')

    def run(self, args):
        if args.url_override:
            url = args.url_override
        else:
            stack = self.get_stack(args.stack)
            if stack is None:
                raise SystemExit('Stack not found: %s' % args.stack)
            url = URL_FORMAT % {"region": stack.region}
        es = ES(url)
        must = [{
            "wildcard": {"containerName": "*%s*" % args.stack}
        }]
        if args.image:
            must.append({"term": {"image_tag": args.image}})
        if args.query:
            must.append({"query_string": {"query": ' '.join(args.query)}})
        if args.level is not None:
            levels = self.levels[self.levels.index(args.level):]
            must.append({"terms": {"priority": levels}})
        rows = list(reversed(es.search(
            'logging-*',
            limit=args.n,
            must=must, sort_order='desc'
        )))
        seen = set(x['sort'][1] for x in rows)
        table_options = self.table_args_parser(args)
        table = self.make_table(args.verbose,
                                args.highlight or [],
                                table_options)
        table.print(rows)
        if not args.follow:
            return
        ts_filter = {"gte": 0}
        must.append({
            "range": {
                "timestamp": ts_filter
            }
        })
        interval = 0.200
        while True:
            if rows:
                last_row_sort = rows[-1]['sort']
                ts_filter['gte'] = last_row_sort[0]
            # Note, this handles paging properly but setting limit to a high
            # value is more performant.
            rows = es.search('logging-*', must=must, limit=10000)
            new = False
            for x in rows:
                if x['sort'][1] not in seen:
                    table.print_row(x)
                    seen.add(x['sort'][1])
                    new = True
            if new:
                interval = 0
            elif interval < 30:
                interval += 0.200
            time.sleep(interval)

    def pretty_ts(self, dt):
        buf = []
        if base.localnow() - dt > datetime.timedelta(hours=12):
            buf.append(base.formatdate(dt))
        buf.append(base.formattime(dt, format='%I:%M:%S.%f %p'))
        return ', '.join(buf)

    def make_table(self, verbose, highlight, table_options):
        colored_levels = {
            "DEBUG": "<dim>DEBUG</dim>",
            "INFO": "<blue>INFO</blue>",
            "WARNING": "<b>WARNING</b>",
            "ERROR": "<red>ERROR</red>",
            "err": "<red>ERROR</red>",
            "CRITICAL": "<blink><b><red>CRITICAL</red></b></blink>",
        }

        def ts_acc(x):
            return self.pretty_ts(base.parse_ts(
                x['_source']['timestamp'] + 'Z'
            ))

        def level_acc(x):
            return colored_levels.get(x['_source']['priority'],
                                      'XXX-%s' % x['_source']['priority'])

        def tag_acc(x):
            return x['_source']['containerTag']

        def msg_acc(r):
            msg = r['_source']['message']
            for hl in highlight:
                msg = msg.replace(hl, '<reverse><b>%s</b></reverse>' % hl)
            return msg

        def hostname_acc(x):
            return x['_source']['hostname'][:19]

        def container_name_acc(x):
            return x['_source']['containerName'][12:24]

        if not verbose:
            config = {
                "headers": [
                    'Time',
                    'Level',
                    'Image',
                    'Message',
                ],
                "columns": [
                    {"minwidth": 13},
                    None,
                    None,
                    None,
                ],
                "accessors": [
                    ts_acc,
                    level_acc,
                    tag_acc,
                    msg_acc,
                ]
            }
        else:
            config = {
                "headers": [
                    'Time',
                    'Level',
                    'Container',
                    'Host',
                    'Message',
                ],
                "columns": [
                    {"minwidth": 13},
                    None,
                    None,
                    None,
                    None,
                ],
                "accessors": [
                    ts_acc,
                    level_acc,
                    container_name_acc,
                    hostname_acc,
                    msg_acc,
                ]
            }
        config = {**config, **table_options}
        return shellish.Table(**config)