Dotfiles_Generator/run.py

#!/usr/bin/env python3

# Recursively read for all files
template_directory = './templates/'

# Written to when generation is done
out_directory = './out/'

# Applied in order, later files override earlier files
env_files = ['./vars.env', './local.env']

import argparse
import difflib
import jinja2
import sys
import re
import os

parser = argparse.ArgumentParser(description='Fill out configuration templates.')
parser.add_argument('-f', '--force', default=False, action=argparse.BooleanOptionalAction, help='Force overwriting of existing files in output')
parser.add_argument('-l', '--link', default=False, action=argparse.BooleanOptionalAction, help='Automatically create symlinks (based on the symlink comment in the templates, must be on first line of template)')
parser.add_argument('-d', '--diff', default=False, action=argparse.BooleanOptionalAction, help='Show the diff of configuration files if file already exists')
parser.add_argument('--dryrun', default=False, action=argparse.BooleanOptionalAction, help='Do not modify any files.')
args = parser.parse_args()
del parser

def main():
    key_value_store = create_key_value_store()
    template_files = find_all_templates()
    out_files = convert_templates_to_out_file(template_files)

    convert_all_templates(template_files, key_value_store, out_files)

def create_key_value_store():
    out = {}

    for env_file in env_files:
        with open(env_file) as file:
            data = file.read()

        for line in data.split('\n'):
            m = re.match(r'^([^={}]+)=([^={}\n]+)', line)
            if m is None:
                continue

            out[m.group(1)] = m.group(2)

    return out

def find_all_templates():
    files = []
    
    dirlist = [template_directory]
    
    while len(dirlist) > 0:
        for (dirpath, dirnames, filenames) in os.walk(dirlist.pop()):
            dirlist.extend(dirnames)
            files.extend(map(lambda n: os.path.join(*n), zip([dirpath] * len(filenames), filenames)))

    return files

def convert_templates_to_out_file(template_files):
    out_files = []

    real_template_dir = os.path.realpath(template_directory)
    real_out_dir = os.path.realpath(out_directory)

    for template_file in template_files:
        real_template_file = os.path.realpath(template_file)
        out_files += [real_template_file.replace(real_template_dir, real_out_dir)]

    return out_files

def convert_all_templates(template_files, key_value_store, out_files):
    env = jinja2.Environment()

    for (template, out) in zip(template_files, out_files):
        convert_template(env, template, key_value_store, out)

def convert_template(env, template, key_value_store, out):
    with open(template, 'r') as file:
        source = file.read()

    template_permissions = os.stat(template).st_mode & 0o777
    default_permissions = 0o666 & ~get_current_umask()

    template = env.from_string(source)
    rendered = template.render(**key_value_store)

    out_dir_name = os.path.dirname(out)
    if not os.path.isdir(out_dir_name):
        dryrun_safe_mkdir(out_dir_name)

    if os.path.isfile(out):
        if is_content_different(out, rendered):
            if args.diff:
                with open(out) as file:
                    current = file.read()

                for line in difflib.unified_diff(current.split('\n'), rendered.split('\n'), fromfile=out, tofile=out+'.new', lineterm=''):
                    print(line, file=sys.stderr)

            if not args.force:
                print(f"File `{out}` already exists, will not overwrite. Rerun with `-f` to force overwriting existing files.", file=sys.stderr)
                return

            dryrun_safe_write(out, rendered)

            if template_permissions != default_permissions:
                dryrun_safe_chmod(out, template_permissions)
        else:
            current_permissions = os.stat(out).st_mode & 0o777
            if template_permissions != current_permissions:
                dryrun_safe_chmod(out, template_permissions)

    if args.link:
        first_line = source.split('\n')[0]
        m = re.search(r'symlink\{([^}]+)\}', first_line)

        if m is None:
            return

        symlink_location = os.path.expanduser(m.group(1))

        create_symlink(out, symlink_location)

def create_symlink(source, destination):
    if os.path.exists(destination):
        if os.path.islink(destination) and os.readlink(destination) == source:
            return

        if not args.force:
            print(f'`{destination}` already exists, will not overwrite. Rerun with `-f` to force overwiring existing files.', file=sys.stderr)
            return

        dryrun_safe_remove(destination)

    dest_dir_name = os.path.dirname(destination)
    if not os.path.isdir(dest_dir_name):
        dryrun_safe_mkdir(dest_dir_name)

    dryrun_safe_create_symlink(source, destination)

def is_content_different(file_location, test_content):
    with open(file_location) as file:
        content = file.read()

    return content != test_content

def get_current_umask():
    current_umask = os.umask(0o022)
    os.umask(current_umask)

    return current_umask

def dryrun_safe_write(location, content):
    if not args.dryrun:
        with open(location, 'w') as file:
            file.write(content)
    else:
        print(f'Would write to `{location}`.')

def dryrun_safe_chmod(location, permissions):
    if not args.dryrun:
        os.chmod(location, permissions)
    else:
        print(f'Would change the permissions of {location} to {oct(permissions)}')

def dryrun_safe_remove(location):
    if not args.dryrun:
        os.remove(location)
    else:
        print(f'Would remove `{location}`.')

def dryrun_safe_mkdir(dir):
    if not args.dryrun:
        os.makedirs(dir)
    else:
        print(f'Would create directory `{dir}`.')

def dryrun_safe_create_symlink(source, destination):
    if not args.dryrun:
        os.symlink(source, destination)
    else:
        print(f'Would create symlink `{destination}`->`{source}`.')

if __name__ == '__main__':
    main()