#!/usr/bin/python

from armory.database.repositories import BaseDomainRepository, UserRepository
from ..ModuleTemplate import ModuleTemplate
import csv
import glob
import six

try:
    from io import BytesIO
except ImportError:
    from StringIO import StringIO as BytesIO


class Module(ModuleTemplate):
    """
    Unfortunately, this is a Windows only deal. So we output a list of commands
    to be run, and have a way of ingesting the output. Since metadata is more
    junk then good data, items have to be manually verified
    """

    name = "PowerMeta"

    def __init__(self, db):
        self.db = db
        self.BaseDomain = BaseDomainRepository(db, self.name)
        self.User = UserRepository(db, self.name)

    def set_options(self):
        super(Module, self).set_options()

        self.options.add_argument(
            "-c",
            "--command",
            help="Print out a list of commands to run",
            action="store_true",
        )
        self.options.add_argument(
            "-ps", "--powershell", help="Filename for a .ps1 command output"
        )
        self.options.add_argument(
            "-i", "--import_file", help="import CSV file generated by PowerMeta"
        )
        self.options.add_argument(
            "-f", "--import_folder", help="import CSV files generated by PowerMeta"
        )

    def run(self, args):
        # pdb.set_trace()

        if args.command:
            res = []
            res.append("Import-Module PowerMeta.ps1")
            res += [
                "Invoke-PowerMeta -TargetDomain %s -Download -Extract -ExtractAllToCsv %s.csv"
                % (b.domain, b.domain)
                for b in self.BaseDomain.all()
            ]

            if args.powershell:
                open(args.powershell, "w").write("\r\n".join(res))
                print("Command saved as %s" % args.powershell)
            else:
                print("\n".join(res))

        if args.import_file:

            self.process_domain(args.import_file, args)

        if args.import_folder:
            if args.import_folder[-1] == "/":
                args.import_folder = args.import_folder[:-1]

            files = glob.glob(args.import_folder + "/*.csv")
            for f in files:
                print("Processing %s" % f)
                self.process_domain(f, args)

    def process_domain(self, csvfile, args):
        domain = csvfile.split("/")[-1].split(".csv")[0]

        domain_obj = self.BaseDomain.find(domain=domain)

        if not domain_obj:
            print("Error: Domain not found: %s" % domain)
            return

        csvdata = open(csvfile).read().decode("utf-16").encode("utf-8")

        csvreader = csv.reader(BytesIO(csvdata), delimiter=",", quotechar='"')
        headers = csvreader.next()
        searchable_headers = ["Author", "Creator"]
        indexes = [headers.index(s) for s in searchable_headers if s in headers]

        data = []
        for row in csvreader:
            for i in indexes:
                data.append(row[i])

        data = list(set(data))  # Dedup

        for d in data:
            if d.strip():
                res = six.input("Is %s a valid name? [y/N] " % d)
                if res and res[0].lower() == "y":
                    if " " in d:
                        if ", " in d:
                            first_name = d.split(", ")[1]
                            last_name = d.split(", ")[0]
                        else:
                            first_name = d.split(" ")[0]
                            last_name = " ".join(d.split(" ")[1:])

                        created, user = self.User.find_or_create(
                            first_name=first_name, last_name=last_name
                        )
                        if created:
                            print("New user created")
                        user.domain = domain_obj
                    # else:
                    #     created, user = self.User.find_or_create(username=d)
                    #     if created:
                    #         print("New username created")
                    #     user.domain = domain_obj
                    user.update()

        self.User.commit()
