#!/usr/bin/env python
import argparse
import os
import subprocess
import sys
import locale

ALLOWED_SCRIPT_EXTENSIONS = ['.py', '.sh']


def run_script(path):
    executable = 'python' if path.endswith('.py') else 'sh'
    ret = subprocess.Popen(
        '{} {}'.format(executable, path), shell=True, stdin=subprocess.PIPE,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True)
    out, err = ret.communicate()
    return out.decode(locale.getpreferredencoding(), errors='replace')


def run_scripts(path):
    output = ""
    for item in sorted(os.listdir(path)):
        item_path = os.path.join(path, item)
        extension = '.{}'.format(item_path.split('.')[-1])
        if os.path.isfile(item_path) and extension in ALLOWED_SCRIPT_EXTENSIONS:
            output += run_script(item_path)
    return output


def generate_motd(source='/etc/motdgen.d/', outfile=None):
    message = run_scripts(source)
    if outfile:
        # Create output dir if missing
        outdir = os.path.dirname(outfile)
        if outdir and not os.path.isdir(outdir):
            os.mkdir(outdir)

        with open(outfile, 'w') as f:
            f.write(message)
    else:
        sys.stdout.write(message)


def cli():
    parser = argparse.ArgumentParser()
    parser.add_argument('-s',
                        '--source',
                        help='Directory containing scripts to generate MOTD',
                        dest='source',
                        default='/etc/motdgen.d/')
    parser.add_argument('-o',
                        '--output',
                        help='File to dump generated MOTD',
                        dest='dest',
                        default=None)
    args = parser.parse_args()
    generate_motd(args.source, args.dest)


if __name__ == '__main__':
    cli()
