#!/usr/share/python2
# -*- coding: utf-8 -*-
#
# wpa_config - a config manager for wpa_supplicant
#
# Author: slowpoke <mail+git@slowpoke.io>
#
# This program is Free Software under the non-terms
# of the Anti-License. Do whatever the fuck you want.
#
# Github: https://github.com/proxypoke/wpa_config

"""wpa_config - a small helper for wpa_supplicant."""

from __future__ import print_function

import argparse
import subprocess
import os


CONFIG_ROOT = "/etc/wpa_config"
CONFIG_SELF = os.path.join(CONFIG_ROOT, "wpa_config.conf")
CONFIG_HEAD = os.path.join(CONFIG_ROOT, "wpa_supplicant.conf.head")
CONFIG_TAIL = os.path.join(CONFIG_ROOT, "wpa_supplicant.conf.tail")
NETWORK_DIR = os.path.join(CONFIG_ROOT, "networks.d")
WPA_SUPPLICANT_CONFIG = "/etc/wpa_supplicant/wpa_supplicant.conf"


class FileExistsError(Exception):
    pass


class Network(object):

    """Represents a single network block in wpa_supplicant.conf."""

    def __init__(self, ssid, **opts):
        self.ssid = ssid
        self.opts = opts

    def __repr__(self):
        string = 'network={\n'
        string += '\tssid="{}"\n'.format(self.ssid)
        for opt, val in self.opts.items():
            string += '\t{}={}\n'.format(opt, val)
        string += '}'
        return string

    def write(self, overwrite=False, network_dir=NETWORK_DIR):
        u"""Write to the appropriate config file.

        Will refuse to overwrite an already present file unless explicitly told
        to do so via the 'overwrite' parameter.

        Arguments:
            overwrite -- ignore any existing files. Defaults to False.
            network_dir -- the directory in which the config files reside.
                Defaults to NETWORK_DIR (global setting).

        Raises:
            FileExistsError -- File exists and overwrite is False.
            IOError -- Permission issue, wrong file type, other OS issue…

        """
        cfile = cfile_path(self.ssid, network_dir)
        if os.path.exists(cfile) and not overwrite:
            raise FileExistsError
        with open(cfile, "w") as f:
            f.write(repr(self))

    @classmethod
    def from_file(cls, ssid, network_dir=NETWORK_DIR):
        """Create a new Network object from a config file."""
        cfile = cfile_path(ssid, network_dir)
        with open(cfile) as f:
            return cls.from_string(f.read())

    @classmethod
    def from_string(cls, string):
        """Create a new network object by parsing a string."""
        lines = string.split("\n")
        # throw away the first and last lines, and remove indentation
        lines = [line.strip() for line in lines[1:-1]]
        opts = {}
        for line in lines:
            split = line.split("=")
            opt = split[0]
            if len(split) == 2:
                value = split[1]
            else:
                value = "=".join(split[1:])
            opts[opt] = value

        # remove the SSID from the other options and strip the quotes from it.
        ssid = opts["ssid"][1:-1]
        del opts["ssid"]

        return cls(ssid, **opts)


def cfile_path(ssid, network_dir=NETWORK_DIR):
    """Return the path of the config file for an SSID."""
    return os.path.join(network_dir, ssid + ".conf")


def merge_cfiles(network_dir=NETWORK_DIR):
    """Merge all config files in the network_dir directory."""
    cfiles = [os.path.join(network_dir, cfile)
              for cfile in os.listdir(network_dir)
              if cfile.endswith(".conf")]
    merge = ""
    for cfile in cfiles:
        with open(cfile) as f:
            merge += f.read()
    return merge


def mkbackup(path):
    """Create a backup file of the given file, named ${FILE}.bkp."""
    with open(path + ".bkp", "w") as bkp_file, open(path) as orig_file:
        bkp_file.write(orig_file.read())


def extract_networks(path):
    """Extract all network blocks from a file.

    Returns a list of Network objects.

    """
    with open(path) as netfile:
        config = netfile.read()
    networks = []
    in_block = False
    netblock = []
    for line in config.split("\n"):
        line = line.strip()
        if line.startswith("#"):
            continue
        # this is really crappy "parsing" but it should work
        if line == "network={" and not in_block:
            in_block = True
        if in_block:
            netblock.append(line)
        if line == "}" and in_block:
            in_block = False
            nb = "\n".join(netblock)
            networks.append(Network.from_string(nb))
            netblock = []
    return networks


# Commands for the CLI frontend

def add(args):
    network = Network(args.ssid)
    if args.open:
        network.opts["key_mgmt"] = "NONE"
    else:
        # check passphrase length
        l = len(args.passphrase)
        if l < 8 or l > 63:
            print("Passphrase must be 8..63 characters.")
            exit(1)
        network.opts["psk"] = '"{}"'.format(args.passphrase)

    if args.print:
        print(network)
        return

    try:
        network.write(network_dir=args.directory,
                      overwrite=args.force)
    except OSError, e:
        print("Couldn't write config file: {}".format(e))
        exit(1)
    except FileExistsError, e:
        print("Config file for '{}' already exists.".format(network.ssid),
              "Refusing to overwrite unless explicitly told to. (-f/--force)")
        exit(1)


def make(args):
    with open(CONFIG_HEAD) as headfile:
        config_head = headfile.read()
    with open(CONFIG_TAIL) as tailfile:
        config_tail = tailfile.read()

    config = "# This file was generated by wpa_config. DO NOT EDIT."
    config += "\n\n"
    config += config_head
    config += "\n\n"
    config += merge_cfiles(args.directory)
    config += "\n\n"
    config += config_tail

    if args.print:
        print(config)
        exit(0)

    try:
        mkbackup(WPA_SUPPLICANT_CONFIG)
        print("Backup of wpa_supplicant.conf written to {}.".format(
            WPA_SUPPLICANT_CONFIG + ".bkp"))
        with open(WPA_SUPPLICANT_CONFIG, "w") as cfile:
            cfile.write(config)
    except OSError, e:
        print("Couldn't write config file: {}".format(e))
        exit(1)


def list_networks(args):
    for cfile in os.listdir(args.directory):
        if cfile.endswith(".conf"):
            print(os.path.splitext(cfile)[0])


def delete(args):
    cfile = cfile_path(args.ssid, args.directory)
    if not os.path.isfile(cfile):
        print(
            "Network '{}' is not configured with wpa_config.".format(cfile))
        exit(1)
    os.unlink(cfile)


def show(args):
    cfile = cfile_path(args.ssid, args.directory)
    if not os.path.isfile(cfile):
        print(
            "Network '{}' is not configured with wpa_config.".format(cfile))
        exit(1)
    with open(cfile) as network:
        print(network.read())


def edit(args):
    editor = os.getenv("EDITOR")
    if editor is None:
        print("You don't appear to have an editor set.")
        exit(1)
    cfile = cfile_path(args.ssid, args.directory)
    if not os.path.isfile(cfile) and not args.create:
        print("Config for '{}' does not exist. Use -c to create it.".format(
            args.ssid))
        exit(1)
    try:
        subprocess.call([editor, cfile])
    except OSError, e:
        print("Failed to invoke editor: {}".format(e))


def migrate(args):
    networks = extract_networks(args.config)
    exit_status = 0
    for network in networks:
        try:
            network.write(network_dir=args.directory,
                          overwrite=args.force)
        except FileExistsError:
            print("Config for '{}' already exists.".format(network.ssid),
                  "Refusing to overwrite unless explicitly",
                  "told to. (-f/--force)")
            exit_status = 1
        except OSError, e:
            print("Couldn't write config file: {}".format(e))
            exit_status = 1
    exit(exit_status)


def main():
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers()

    def ssid_opt(parser):
        parser.add_argument("ssid", type=str, help="network name")

    def dir_opt(parser):
        parser.add_argument(
            "-d", "--directory", type=str, default=NETWORK_DIR,
            help="directory to use as base (default: {})".format(NETWORK_DIR))

    add_mode = commands.add_parser("add", help="add a network")

    ssid_opt(add_mode)
    add_mode.add_argument("passphrase", type=str, help="wpa passphrase",
                          nargs="?", default="")
    dir_opt(add_mode)
    add_mode.add_argument("-p", "--print", action="store_true",
                          help="print config instead of writing it")
    add_mode.add_argument("-f", "--force", action="store_true",
                          help="overwrite existing config file")

    netopts = add_mode.add_argument_group("network options")
    netopts.add_argument("-o", "--open", action="store_true",
                         help="configure for an open network")

    add_mode.set_defaults(func=add)

    delete_mode = commands.add_parser("del", help="remove a network")
    ssid_opt(delete_mode)
    dir_opt(delete_mode)
    delete_mode.set_defaults(func=delete)

    show_mode = commands.add_parser("show", help="show network configuration")
    ssid_opt(show_mode)
    dir_opt(show_mode)
    show_mode.set_defaults(func=show)

    edit_mode = commands.add_parser(
        "edit", help="edit a network with an editor")
    ssid_opt(edit_mode)
    dir_opt(edit_mode)
    edit_mode.add_argument("-c", "--create", action="store_true",
                           help="create the config file if it doesn't exist")
    edit_mode.set_defaults(func=edit)

    list_mode = commands.add_parser("list", help="list configured networks")
    dir_opt(list_mode)
    list_mode.set_defaults(func=list_networks)

    make_mode = commands.add_parser("make", help="create the config")
    dir_opt(make_mode)
    make_mode.add_argument("-p", "--print", action="store_true",
                           help="print config instead of writing it")
    make_mode.set_defaults(func=make)

    migrate_mode = commands.add_parser("migrate",
                                       help="migrate your old config")
    dir_opt(migrate_mode)
    migrate_mode.add_argument("-c", "--config", type=str,
                              help="path to wpa_supplicant.conf",
                              default=WPA_SUPPLICANT_CONFIG)
    migrate_mode.add_argument("-f", "--force", action="store_true",
                              help="overwrite existing config files. " +
                              "(DANGEROUS)")
    migrate_mode.set_defaults(func=migrate)

    help_mode = commands.add_parser("help", help="print help for a mode")
    help_mode.set_defaults(func=lambda _: parser.print_help())

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
