data2elk.py

#!/usr/bin/env python
"""data2elk.py feeds csv data into Elasticsearch-Logstash-Kibana"""

import os
import sys
import csv
import glob
from elk_config import ConfigElk as config_elk

CONFIG_TEMPLATE = """# logstash config, generated by data2elk.py
input {
  file {
    path => "%s"
    type => "csv"
    start_position => "beginning"
    %s
  }
}
 
filter { 
  csv {
     columns => %s
     separator => "%s"
  } 
}
 
output {
  elasticsearch { 
    action => "index"
    hosts => "localhost" 
    index => "logstash-%%{+YYYY.MM.dd}"
    workers => 1
  }
  stdout { codec => rubydebug }
}
"""

def get_columns(path, delimiter=',', quotechar='"'):
    """Extract column names from the CSV

    @param path: path to CSV file
    @param delimiter: CSV delimiter character
    @param quotechar: character used to quote fields in CSV
    @return: list of column names
    """
    with open(path, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=delimiter, quotechar=quotechar)
        header = reader.next()  

    return header


def generate_config(config_path, csv_path, columns, ignore_sincedb=False, 
                    delimiter=',', quotechar='"'):
    """Generate Logstash config file.

    @param config_path: path to write Logstash config to.
    @param csv_path: path to csv input data
    @param columns: column names in the csv
    @param ignore_sincedb: set True to ingest everything, 
        False to ingest only new data 
    @param delimiter: CSV delimiter character
    @param quotechar: character used to quote fields in CSV
    @return: None
    """
    config_dir = os.path.dirname(config_path)
    if not os.path.exists(config_dir):
        os.makedirs(config_dir)

    # if the specified path is a dir, add a wildcard
    path = os.path.abspath(csv_path)
    if os.path.isdir(path):
        path = os.path.join(path, '*')

    ignore_sincedb = 'sincedb_path => "/dev/null"' if ignore_sincedb else ''
    data = CONFIG_TEMPLATE % (path, ignore_sincedb, columns, delimiter)
    with open(config_path, 'w') as fd:
        fd.write(data)


def which(program, all=False):
    """Locates the program in the directories specified by the PATH environment variable.

    @param program: the program to locate
    @param all: boolean, specify True to find all instances of the program found, 
                otherwise just the first is returned.
    @return: The path found, or None.  If all is True, returns list of all paths found.
    """              
    found = []
    for path in os.getenv("PATH").split(os.path.pathsep):
        full_path = os.path.join(path, program)
        if os.path.exists(full_path):
            if all:
                found.append(full_path)
            else:
                return full_path

    if all:
        return found


def _latest_file(dir):
    """Return absolute path to latest file in directory dir."""
    return max(glob.glob(os.path.join(dir, '*')), key=os.path.getctime)


def run(path, delimiter=',', quotechar='"', logstash_config='/etc/logstash/conf.d/logstash.conf',
        restart_logstash=True, ip='localhost', skip_install=True, ignore_sincedb=False):
    if not os.path.exists(path):
        sys.stdout.write("File does not exist: {}\n".format(path))
        sys.exit(1)  

    if sys.platform.startswith('linux') and not args.skip_install:
        config_elk(args.ip)

    # if a dir was specified, pick the latest file in the dir to get the
    # header from, otherwise use the specified file
    f = _latest_file(args.file) if os.path.isdir(args.file) else args.file
    columns = get_columns(f, delimiter=args.delimiter, quotechar=args.quotechar) 

    generate_config(args.output, args.file, columns, ignore_sincedb=ignore_sincedb,
                    delimiter=args.delimiter, quotechar=args.quotechar)

    LOGSTASH = "/usr/share/logstash/bin/logstash" # which('logstash')
    LOGSTASH_SETTINGS_DIR = "/etc/logstash"
    if not os.path.exists(LOGSTASH):
        sys.stdout.write("No logstash installation found\n")
        sys.exit(1)    

    # check if generated config is valid, raises CalledProcessError is raised if config is invalid.
    # bin/logstash -t -f /etc/logstash/logstash.conf 
    with open(os.devnull, 'w') as devnull_fd: # used to suppress output
        subprocess.check_call([LOGSTASH, "--config.test_and_exit", 
                               "--path.config", logstash_config,
                               "--path.settings", LOGSTASH_SETTINGS_DIR], 
                              stdout=devnull_fd)
  
    if restart_logstash:
        subprocess.call([LOGSTASH, '-f', args.output])            


if __name__ == "__main__":
    import argparse
    import subprocess


    DEFAULT_CONFIG_PATH = '/etc/logstash/conf.d/logstash.conf'

    parser = argparse.ArgumentParser(description="Process data for Elasticsearch/Logstash/Kibana")
    parser.add_argument('-f', '--file', metavar='FILE', help='path to csv input or dir')
    parser.add_argument('--delimiter', default=',',
                        help="""csv delimiter character, ',' by default""")
    parser.add_argument('--quotechar', default='"',
                        help="""csv quote character, '"' by default""")
    parser.add_argument('-o', '--output', metavar='FILE', default='/etc/logstash/logstash.conf',
                        help='path to logstash config output, defaults to {}'.format(DEFAULT_CONFIG_PATH))
    parser.add_argument('-r', '--restart-logstash', action='store_true', default=False,
                        help='restart logstash with the generated config, if it is not already running as a daemon')
    parser.add_argument('-i', '--ip', default='localhost',
                        help='IP address for Kibana and ElasticSearch instance, defaults to localhost.')
    parser.add_argument('--skip-install', action='store_true', default=False,
			            help='skip ELK installation')
    parser.add_argument('--ignore-sincedb', action='store_true', default=False,
                        help='ignore sincedb to ingest previously ingested data')
    args = parser.parse_args()    

    if not args.file:
        sys.stdout.write("No data path specified\n")
        sys.exit(1)

    run(args.file, delimiter=args.delimiter, quotechar=args.quotechar, 
        logstash_config=args.output, restart_logstash=args.restart_logstash,
        ip=args.ip, skip_install=args.skip_install, ignore_sincedb=args.ignore_sincedb)