# -*- coding: utf-8 -*-
"""safire.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Op2TEgznirfS62ewpEP8XfXUrElNpYs6

### Safire

###### Instructions

1. If using Colab:
  - Save copy of the ipynb file to your drive (Colab Notebooks folder)
  - Run the init code cell to auth to your My Drive account. This gives you access to Drive where you will store your creds, svcaccts and xlsx/csv files.

2. If Running locally then change the config to local folders
3. Run the Config cell once to create empty /creds, /data and /svcaccts folders in the /safire folder if they do not exist.
4. If you have creds.json and token files, drag them to the /safire/creds you just created.
5. If you do not have creds.json and token files then create a creds.json file and drop it in safire/creds. Then run `auth projects` and `auth groups` (see [github](https://github.com/88lex/safire) for details.

^^ All of this you do once, after which you will have a structure that looks like this:
```
/content/drive/My Drive
  /safire
    /creds
      /creds.json
      /token.pickle
      /grptoken.pickle
    /data
    /svcaccts
```

After creating your credentials and folders:
6. Enter one or more commands in the Commands section (empty cells are ignored). Command syntax is:
  - `list projects`
  - `list all`
  - `list counts` will show you item counts for projects, drives, groups, etc.
  - `add projects` will add the number of projects using a prefix that you entered in the config file
  - `add projects 5 1 myproj` will create 5 projects starting with number 0001 using a prefix of myproj. Result: myproj0001, myproj0002 .. myproj0005

7. Run commands using `Runtime - Run All` or `Ctrl-F9`
  - NOTE: If you run only the Command cell it will do nothing. You need to Run All in order to run all of the code.

###### **Init**

- Init needs to be run once each time the Colab instance or Jupyter notebook is started
"""

import os, importlib, sys, uuid, pickle
import pandas as pd
from base64 import b64decode
from glob import glob
from json import loads
from time import sleep
from pathlib import Path
from google.auth.transport.requests import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from shutil import copyfile

if importlib.util.find_spec('fire') is None: os.system('pip install fire')
import fire

#@markdown Root dir when using Colab GDrive mount(recommended to leave default):
try:
    from google.colab import drive
    gdrive_mount = '/content/drive' #@param ['/content/drive'] {allow-input: true}
    gd_output_dir = '/My Drive/safire' #@param ['/My Drive/safire'] {allow-input: true}
    drive.mount(gdrive_mount)
    safire_dir = gdrive_mount + gd_output_dir
except:
    base_path = str(Path.home())
    safire_dir = os.path.join(base_path, "safire")
    os.makedirs(safire_dir, exist_ok=True)
    cfg_file = os.path.join(safire_dir, "config.py")
    if os.path.exists(cfg_file):
        copyfile(cfg_file, "config.py")
    else:
        copyfile("default_config.py", "config.py")
        copyfile("default_config.py", cfg_file)

    ## @markdown Root directory when safire running locally (dropdown choices):
    # local_safire_dir = '~/safire' #@param ['~/safire', '/opt/safire/safire', '~/../opt/safire/safire', "safire"] {allow-input: true}
    # safire_dir = local_safire_dir
# @markdown External config file located in local_safire_dir. Default: config.py:
use_external_config = True  # @param {type:"boolean"}
ext_conf_file = "config.py"  # @param ["config.py"] {allow-input: true}
ext_conf_file = os.path.join(safire_dir, ext_conf_file)


# except:
#     #@markdown Root directory when safire running locally (dropdown choices):
#     local_safire_dir = '~/safire' #@param ['~/safire', '/opt/safire/safire', '~/../opt/safire/safire', "safire"] {allow-input: true}
#     safire_dir = local_safire_dir
# #@markdown External config file located in local_safire_dir. Default: config.py:
# use_external_config = True #@param {type:"boolean"}
# ext_conf_file = "/config.py" #@param ["/config.py"] {allow-input: true}
# ext_conf_file = safire_dir + ext_conf_file

"""###### Config
- Tells safire where to read and write files, how many 'things' to create and what to call them
"""

import os
import random
import string
from pathlib import Path
from time import sleep

#@markdown Project prefix must be unique across all of Google. If blank a random project prefix will be generated
#@markdown :  NOTE - DO NOT keep sample name below = `project`
project_prefix = "project"        #@param [ ""] {allow-input: true}
if not project_prefix:
    letters = string.ascii_lowercase
    project_prefix = ''.join(random.choice(letters) for i in range(10))
    print(f"Project prefix = {project_prefix}")

#@markdown Prefix for service account emails is =required= (can be anything you like)
email_prefix = "sa" #@param [ "sa", "svcacc", "bananas"] {allow-input: true}

#@markdown Prefix for JSON keys is =optional=
json_prefix = ""          #@param {type:"string"}
#@markdown Enter a group email address if using group functions (e.g. mygroup@domain.com)
group_address = ""        #@param [ "", "mygroup@domain.com"] {allow-input: true}

#@markdown Slide bars and dropdowns can be adjusted with arrow keys
num_new_projects = 1      #@param {type:"slider", min:1, max:200, step:1}
sas_per_project = 100     #@param {type:"slider", min:1, max:100, step:1}
# max_projects =  100     #@param {type:"integer"} 
next_project_num = 1      #@param {type:"integer"}
next_sa_num = 1           #@param {type:"integer"}
#@markdown Note: JSON numbering can be independent of SA numbering
next_json_num = 1         #@param {type:"integer"}

#@markdown ---

# #@markdown Enter 'safire_dir'where creds, SAs and data files will be saved. Default is 
# #@markdown `/content/drive/My Drive/safire`
# safire_dir = "/content/drive/My Drive/safire"#@param ["/content/drive/My Drive/safire"] {allow-input: true}

#@markdown These are subfolders of 'safire_dir' above. 
#@markdown Creds and tokens are stored in the same folder
cred_folder = "/creds"       #@param {type:"string"}
cred_path = safire_dir + cred_folder
cred_file = "/creds.json"    #@param {type:"string"}
credentials = cred_path + cred_file
token_file = "/token.pickle" #@param {type:"string"}
token = f"{cred_path}/token.pickle"
sa_folder = "/svcaccts"      #@param {type:"string"}
sa_path = safire_dir + sa_folder
data_folder = "/data"        #@param {type:"string"}
data_path = safire_dir + data_folder

#@markdown Create and store a separate token for groups. You can use the same creds
group_cred_file = "/creds.json" #@param {type:"string"}
group_credentials = cred_path + group_cred_file
group_token_file = "/grptoken.pickle" #@param {type:"string"}
group_token = cred_path + group_token_file

#@markdown ___

#@markdown Set zero padding for names. e.g. jpad=6 yields 000001.json. Set to 1 if no prefix needed.
ppad = 4             #@param {type:"integer"}
spad = 6             #@param {type:"integer"}
jpad = 6             #@param {type:"integer"}
#@markdown retry and sleep_time help safire deal with slow google apis. Adjust as needed
retry =            5#@param {type:"integer"}
sleep_time =        5#@param {type:"integer"}

## #@markdown [API Service Names and Versions]
DRIVE = ["drive", "v3"]
CLOUD = ["cloudresourcemanager", "v1"]
IAM = ["iam", "v1"]
ADMIN = ["admin", "directory_v1"]
SVCUSAGE = ["serviceusage", "v1"]
SHEETS = ["sheets", "v4"]
svcs_to_enable = ["iam", "drive"]

### Generate empty directories if they do not exist
if not os.path.isdir(safire_dir):
    os.makedirs(safire_dir, exist_ok=True)
    sleep(sleep_time)
if not os.path.isdir(sa_path):
    os.makedirs(sa_path, exist_ok=True)
if not os.path.isdir(cred_path):
    os.makedirs(cred_path, exist_ok=True)
if not os.path.isdir(data_path):
    os.makedirs(data_path, exist_ok=True)

if use_external_config:
    try:
        exec(open(ext_conf_file).read())
    except:
        print(f"Config file = {ext_conf_file} not found, not executed")

"""###### Utils

This section conatins the utils functions then main commands/subcommands
"""

import os
import pickle
import uuid
import pandas as pd

from glob import glob
from pathlib import Path
from time import sleep

from google.auth.transport.requests import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build


class Help:
    """These small functions support repeated activities in other classes/functions"""

    def __init__(self):
        super(Help, self).__init__()

    def _svc(self, service, version, tkn):
        try:
            with open(tkn, "rb") as t:
                creds = pickle.load(t)
            return build(service, version, credentials=creds)
        except:
            print("No valid token found")
            auth = Auth()
            auth.check()

    def _export(self, dset, filter, file_tag, fields, ltype, prt=True):
        df = pd.DataFrame(dset)
        fname = f"{ltype}_list_{filter}_{file_tag}"
        pd.set_option('display.max_rows', None)
        df.to_csv("{}/{}.csv".format(data_path, fname))
        df.to_excel("{}/{}.xlsx".format(data_path, fname))
        if prt:
            print(df[fields])
            print(f"  {len(df)} {ltype} found")
            print(f"\nData for {ltype} exported to {fname}.csv and .xlsx in folder /{data_path}\n")

    def _print1(self, plist, ltype):
        print('', *plist, sep='\n')
        print(f"  {len(plist)} {ltype} found")

    def _pre_pad(self, prefix, pad, number):
        return prefix + (('0' * pad) + str(number))[-pad:]


class BatchJob():

    def __init__(self, service):
        self.batch = service.new_batch_http_request(callback=self.callback_handler)
        self.batch_resp = []

    def add(self, to_add, request_id=None):
        self.batch.add(to_add, request_id=request_id)

    def callback_handler(self, rid, resp, exception):
        response = {'request_id': rid, 'exception': None, 'response': None}
        if exception is not None:
            response['exception'] = exception
        else:
            response['response'] = resp
        self.batch_resp.append(response)

    def execute(self):
        try:
            self.batch.execute()
        except:
            pass
        return self.batch_resp


class Auth:
    """Authorize the app to access your projects, SAs, drives and groups. To generate creds.json go to
    https://developers.google.com/apps-script/api/quickstart/python , click Enable then download a json,
    rename it to creds.json and put a copy in the /creds folder"""

    def __init__(self):
        super(Auth, self).__init__()
        self.scopes_proj = [
            "https://www.googleapis.com/auth/drive",
            "https://www.googleapis.com/auth/cloud-platform",
            "https://www.googleapis.com/auth/iam",
        ]
            # "https://www.googleapis.com/auth/spreadsheets",

        self.scopes_group = [
            "https://www.googleapis.com/auth/admin.directory.group",
            "https://www.googleapis.com/auth/admin.directory.group.member",
        ]
        self.scopes_all = self.scopes_proj + self.scopes_group

    def ask(self):
        pass

    def check(self):
        filelist = [credentials, token, group_credentials, group_token]
        file_exists = [os.path.isfile(i) for i in filelist]
        [print(f"File = {i[0]} exists = {i[1]}") for i in zip(filelist, file_exists)]
        if not file_exists[0]:
            print(
                f"Credentials file is missing. Download from Google console and run 'auth'"
            )
        if not file_exists[2]:
            print(
                f"Group credentials file is missing. Download from Google console and run 'auth'"
                f"Note that the group credentials file is normally the same as the main projects credentials"
                f"file - But you can optionally use separate credentials files. Specify in config.py"
            )
        yes_no = input("Generate token for projects, SAs, drives? y/n: ")
        if yes_no.lower() in ["y", "yes"]:
            self.projects(credentials, token)
        gyes_no = input("Generate token for groups? y/n: ")
        if gyes_no.lower() in ["y", "yes"]:
            self.groups(group_credentials, group_token)
        exit()

    def projects(self, credentials=credentials, token=token):
        """Create an auth token for accessing and changing projects,
        service accounts, json keys and drives"""
        self.get_token(credentials, self.scopes_proj, token)

    def groups(self, credentials=credentials, group_token=group_token):
        """Create an auth token for adding/removing group members"""
        self.get_token(credentials, self.scopes_group, group_token)

    def all(self, credentials=credentials, token=token):
        """Create an auth token with APIs enabled for projects and groups"""
        self.projects(credentials, token)
        self.groups(group_credentials, group_token)

    def get_token(self, credentials, scopes, token):
        if not credentials or os.path.isfile(credentials[0]):
            print(
                f"Credentials file is missing. Download from Google console and save as {credentials}"
            )
            exit()
        cred_file = glob(credentials)
        creds = None
        if os.path.exists(token):
            with open(token, "rb") as tkn:
                creds = pickle.load(tkn)
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                creds.refresh(Request())
            else:
                flow = InstalledAppFlow.from_client_secrets_file(cred_file[0], scopes)
                yes_no = input(
                    "Run local server with browser? If no, generate console link. y/n: "
                )
                if yes_no.lower() in ["y", "yes"]:
                    creds = flow.run_local_server()
                else:
                    creds = flow.run_console()
            with open(token, "wb") as tkn:
                print("Writing/updating token")
                pickle.dump(creds, tkn)
        else:
            print("Credentials and token exist and appear to be valid")
            print(f"credentials = {cred_file[0]} and token = {token}")
            yes_no = input(
                f"Do you want to delete and regenerate your token file = {token}? y/n: "
            )
            if yes_no.lower() in ["y", "yes"]:
                os.remove(token)
                self.get_token(credentials, scopes, token)
            else:
                print("Finished without creating new token")


class Link():
    """Create a symlink between safire's directories and another location"""

    def dirs(self):
        cwd = os.path.dirname(os.path.realpath(__file__))
        dest = f"{str(Path.home())}/safire"
        dest1 = input(f"Choose dir to link. To keep default = {dest} press Enter: ")
        if dest1:
            dest = dest1
        if os.path.exists(dest):
            print(f"Directory {dest} already exists. Exiting.")
            exit()
        print(f"Linking {cwd} to {dest}")
        if os.path.isdir(dest):
            os.unlink(dest)
        os.symlink(cwd, dest)


class Commands:
    """safire creates projects, service accounts(SAs), json keys for SAs and adds SAs as group members.\n Type -h
    or --help after any command for info.  Usage: './safire list projects' or 'safire add members mygroup@domain.com'
    Most commands accept a 'filter'to process subsets. If no filter all items are listed."""

    def __init__(self):
        # Init()
        self.list = List()
        self.add = Add()
        self.remove = Remove()
        self.auth = Auth()
        self.rename = Rename()
        self.link = Link()
        # alias commands for ease of use. e.g. 'safire add projects' = 'safire create projects'
        self.create = self.add
        self.enable = self.add
        self.delete = self.remove

"""###### List"""

class List(Help):
    """List drives, projects, service accounts (SAs), SA json keys, groups and group members. 
    In most cases a filter can be applied."""

    def __init__(self):
        super(List, self).__init__()

    def drives(self, filter="", file_tag="", token=token, prt=True):
        """List team/shared drives. Match 'filter'"""
        drivelist = []
        resp = {"nextPageToken": None}
        drive = self._svc(*DRIVE, token)
        while "nextPageToken" in resp:
            resp = (
                drive.drives()
                .list(
                    fields="nextPageToken, drives(id,name)",
                    pageSize=100,
                    pageToken=resp["nextPageToken"],
                )
                .execute()
            )
            drivelist += resp["drives"]
        drivelist = [i for i in drivelist if str(filter) in i["name"]]
        if not prt:
            return drivelist
        self._export(drivelist, filter, file_tag, ["id", "name"], "drives", prt)

    def projects(self, filter="", file_tag="", token=token, prt=True, exact=False):
        """List projects. Match 'filter'"""
        cloud = self._svc(*CLOUD, token)
        plist = cloud.projects().list().execute()["projects"]
        plist = [i for i in plist if str(filter) in i["projectId"]]
        if exact:
            plist = [i for i in plist if str(filter) == i["projectId"]]
        if not prt:
            return [i["projectId"] for i in plist]
        self._export(
            plist,
            filter,
            file_tag,
            ["projectNumber", "projectId", "name"],
            "projects",
            prt,
        )

    def sas(self, filter="", exact=False, file_tag="", token=token, prt=True):
        """List service accounts for projects. Projects match 'filter'.
        Use '--exact=True' flag for exact project match."""
        svcacctlist, sa_summary = [], []
        projId_list = self.projects(filter, file_tag, token, False, exact)
        # if exact:
        #     projId_list = [i for i in projId_list if i == filter]
        iam = self._svc(*IAM, token)
        for project in sorted(projId_list):
            try:
                resp = (
                    iam.projects()
                    .serviceAccounts()
                    .list(name="projects/" + project, pageSize=100)
                    .execute()
                )["accounts"]
                svcacctlist += resp
                sa_summary += [
                    f"  {len(resp)} service accounts (SAs) found in project: {project}"
                ]
            except:
                sa_summary += [
                    f"  No service accounts (SAs) found in project: {project}"
                ]
                continue
        svcacctlist.sort(key=lambda x: x.get("email"))
        if not prt:
            return (
                [i["email"] for i in svcacctlist],
                [i["uniqueId"] for i in svcacctlist],
            )
        if svcacctlist:
            self._export(svcacctlist, filter, file_tag, ["email"], "svc_accts")
        print("\nSummary:", *sa_summary, sep="\n")

    def groups(self, filter="", file_tag="", group_token=group_token, prt=True):
        """List groups in the authorized account. Match 'filter'"""
        svc = self._svc(*ADMIN, group_token)
        grouplist = svc.groups().list(customer="my_customer").execute()["groups"]
        grouplist = [i for i in grouplist if str(filter) in i["email"]]
        if not prt:
            return [i["email"] for i in grouplist]
        self._export(
            grouplist,
            filter,
            file_tag,
            ["id", "email", "directMembersCount"],
            "groups",
            prt,
        )

    def members(self, filter="", file_tag="", group_token=group_token, prt=True):
        """List members in groups. Groups match 'filter'"""
        memberslist, member_summary = [], []
        group_list = self.groups(filter, file_tag, group_token, False)
        svc = self._svc(*ADMIN, group_token)
        for group in group_list:
            try:
                response = []
                resp = {"nextPageToken": None}
                while "nextPageToken" in resp:
                    resp = (
                        svc.members()
                        .list(groupKey=group, pageToken=resp["nextPageToken"])
                        .execute()
                    )
                    response += resp["members"]
                memberslist += response
            except:
                member_summary += [f"  No members found in group: {group}"]
                continue
            member_summary += [f"  {len(response)} members found in group: {group}"]
        if not prt:
            return sorted([i["email"] for i in memberslist if i["role"] != "OWNER"])
        self._export(memberslist, filter, file_tag, ["email", "role"], "members")
        print("\nSummary:", *member_summary, sep="\n")

    def jsons(self, sa_path=sa_path, filter="", file_tag="", prt=True):
        """alias: jsons = keys. List service account jsons/keys in the svcaccts folder. Match 'filter'"""
        keys = glob("%s/*.json" % sa_path)
        json_keys = []
        [json_keys.append(loads(open(key, "r").read())) for key in keys]
        if filter:
            json_keys = [key for key in json_keys if str(filter) in key["project_id"]]
        if prt and json_keys:
            self._export(
                json_keys, filter, file_tag, ["client_email"], "json_keys", prt
            )
            print(f"There are {len(json_keys)} json keys in {sa_path}")
        elif json_keys:
            return [i["client_email"] for i in json_keys]
        # else:
        #     print(f"No json keys in path: {sa_path} matching filter: {filter}")

    keys = jsons

    def all(self):
        """List all drives, projects, service accounts, json keys, groups and group members.
        Also exports these lists with full data fields to csv and xlsx files in data_path folder"""
        self.drives()
        self.projects()
        self.sas()
        self.groups()
        self.members()
        self.jsons()

    def count(self):
        """Count drives, projects, service accounts, json keys, groups and group members."""
        list_drives = []
        list_drives = self.drives("", "", token, False)
        print("Drive count  :", len(list_drives))
        list_projects = []
        list_projects = self.projects("", "", token, False)
        print("Project count:", len(list_projects))
        list_groups = []
        list_groups = self.groups("", "", group_token, False)
        print("Group count  :", len(list_groups))
        try:
          list_jsons = self.jsons(sa_path, "", "", False)
          print("JSON count   :", len(list_jsons))
        except:
          print("JSON count   : 0")
        list_sas = []
        list_sas, _ = self.sas("", False, "", token, False)
        print("SA count     :", len(list_sas))
        list_members = []
        list_members = self.members("", "", group_token, False)
        print("Member count :", len(list_members))

"""###### Add"""

class Add(Help):
    """Add projects, drives, service accounts(SAs), SA keys/jsons and group members"""

    def __init__(self):
        super(Add, self).__init__()
        self._list = List()

    def projects(
        self,
        num_new_projects=num_new_projects,
        next_project_num=next_project_num,
        project_prefix=project_prefix,
        retry=5,
        token=token,
        ppad=4,
        prt=False,
    ):
        """Create projects in authorized account. Usage: 'safire add projects 1'.
        Uses defaults in config if none specified."""
        cloud = self._svc(*CLOUD, token)
        batch = BatchJob(cloud)
        start_projs = self._list.projects("", "", token, prt)
        start_proj_count = len(start_projs)
        num_proj_created = 0
        curr_projects = []
        proj_count = 0
        print(f"Creating {num_new_projects} new projects")
        retries = retry
        while num_new_projects and retries:
            new_projs = [
                self._pre_pad(project_prefix, ppad, next_project_num + i)
                for i in range(num_new_projects)
            ]
            for proj in new_projs:
                batch.add(
                    cloud.projects().create(body={"project_id": proj, "name": proj})
                )
            batch.execute()
            retries -= 1
            curr_projects = self._list.projects("", "", token, prt)
            proj_count = len(curr_projects)
            num_proj_created = proj_count - start_proj_count
            next_project_num = next_project_num + num_proj_created
            num_new_projects = num_new_projects - num_proj_created
        if num_proj_created < 1:
            print(
                "0 projects created. Your project prefix + number may be used already",
                "somewhere else in Google. Try changing your prefix.",
            )
            exit()
        new_projs = [i for i in curr_projects if i not in start_projs]
        print(f"\nCreated {num_proj_created} projects:", *new_projs, sep="\n")
        print(f"Total project count = {proj_count}")
        print(
            "Sleep briefly to allow backend to register projects. Then enabling services.\n"
        )
        sleep(sleep_time)
        for proj in new_projs:
            self.apis(proj)

    def apis(
        self, filter="", svcs_to_enable=svcs_to_enable, token=token, prt=False
    ):
        """Enables apis for projects. 'drive' and 'iam' apis by default. Automatic when projects are
        created but can be run manually also."""
        svcusage = self._svc(*SVCUSAGE, token)
        projId_list = self._list.projects(filter, "", token, prt)
        batch = BatchJob(svcusage)
        svcs_to_enable = [i + ".googleapis.com" for i in svcs_to_enable]
        for project in projId_list:
            for svc1 in svcs_to_enable:
                print(f"Enabling service: {svc1} in project: {project}")
                batch.add(
                    svcusage.services().enable(
                        name=f"projects/{project}/services/{svc1}"
                    )
                )
            batch.execute()

    def sas(
        self,
        filter="",
        sas_per_project=sas_per_project,
        email_prefix=email_prefix,
        next_sa_num=next_sa_num,
        retry=retry,
        exact=False,
        file_tag="",
        prt=False,
    ):
        """Create N service accounts/SAs in projects which match 'filter'. Usage: 'safire add sas xyz 5'
         will add 5 SAs to all projects containing 'xys' if fewer than 100 exist. Will not overwrite SAs."""
        iam = self._svc(*IAM, token)
        projId_list = self._list.projects(filter, file_tag, token, prt)
        all_sas = []
        for project in projId_list:
            # batch = BatchJob(iam)
            sa_emails, _ = self._list.sas(project, True, file_tag, token, False)
            start_sa_emails = sa_emails
            start_sa_count = len(sa_emails)
            all_sas = all_sas + sa_emails
            count = min(sas_per_project, 100 - len(sa_emails))
            print(
                len(sa_emails), "SAs exist. Creating", count, "SAs in project", project
            )
            new_sas = []
            retries = retry
            while count and retries:
                batch = BatchJob(iam)
                for _ in range(count):
                    while [s for s in all_sas if str(next_sa_num) in s.split("@")[0]]:
                        next_sa_num += 1
                    sa_id = self._pre_pad(email_prefix, spad, next_sa_num)
                    new_sas.append(sa_id)
                    next_sa_num += 1
                    name = f"projects/{project}"
                    body = {
                        "accountId": sa_id,
                        "serviceAccount": {"displayName": sa_id},
                    }
                    batch.add(
                        iam.projects().serviceAccounts().create(name=name, body=body)
                    )
                batch.execute()
                retries -= 1
                sa_emails, _ = self._list.sas(project, True, file_tag, token, False)
                curr_sa_count = len(sa_emails)
                count = count - curr_sa_count + start_sa_count
                sleep(sleep_time/5)
            new_sa_emails = [i for i in sa_emails if i not in start_sa_emails]
            num_sas_created = len(new_sa_emails)
            print(
                f"\nCreated {num_sas_created} sas in {project}:",
                *new_sa_emails,
                sep="\n",
            )
            print(f"Total SAs in {project} count = {len(sa_emails)}")

    def jsons(
        self,
        filter="",
        sa_path=sa_path,
        next_json_num=next_json_num,
        json_prefix=json_prefix,
        jpad=jpad,
        retry=retry,
        file_tag="",
        token=token,
        prt=False,
    ):
        """Create and download json/key files to svcaccts folder. Add to TDs and/or groups."""
        iam = self._svc(*IAM, token)
        projId_list = self._list.projects(filter, file_tag, token, prt)
        for project in sorted(projId_list):
            batch = BatchJob(iam)
            _, sa_uniqueId = self._list.sas(project, True, file_tag, token, prt)
            num_sas = len(sa_uniqueId)
            print(f"Downloading {str(len(sa_uniqueId))} SA keys in project {project}")
            resp = []
            retries = retry
            while len(resp) < num_sas and retries:
                for sa in sa_uniqueId:
                    batch.add(
                        iam.projects()
                        .serviceAccounts()
                        .keys()
                        .create(
                            name=f"projects/{project}/serviceAccounts/{sa}",
                            body={
                                "privateKeyType": "TYPE_GOOGLE_CREDENTIALS_FILE",
                                "keyAlgorithm": "KEY_ALG_RSA_2048",
                            },
                        )
                    )
                resp = [i["response"] for i in batch.execute()]
                retries -= 1
            for i in resp:
                if i is not None:
                    k = (
                        i["name"][i["name"].rfind("/") :],
                        b64decode(i["privateKeyData"]).decode("utf-8"),
                    )[1]
                    json_name = self._pre_pad(json_prefix, jpad, next_json_num)
                    with open("%s/%s.json" % (sa_path, json_name), "w+") as f:
                        f.write(k)
                    next_json_num += 1
            sleep(sleep_time)

    keys = jsons

    def drive(self, td_name):
        """Create a team/shared drive. Usage: 'safire add drive some_name'"""
        drive = self._svc(*DRIVE, token)
        body = {"name": td_name}
        driveId = (
            drive.drives()
            .create(body=body, requestId=str(uuid.uuid4()), fields="id,name")
            .execute()
            .get("id")
        )
        print(f"  Drive ID for {td_name} is {driveId}")
        return td_name, driveId

    def drives(self, *td_list):
        """Create team/shared drives. Usage: 'safire add drives some_filename' containing 
        a list of drive names OR `safire add drives mytd1 mytd2 mytd3`"""
        td_file = ' '.join(sys.argv[3:]).strip('\"')
        if os.path.isfile(td_file):
            with open(td_file, "r") as td_lst:
              td_list = [i for i in td_lst]
        for td_name in td_list:
            print(f"Creating {td_name}")
            self.drive(td_name.rstrip())

    def members(
        self,
        project_filter,
        group_filter,
        retry=5,
        file_tag="",
        group_token=group_token,
        prt=False,
    ):
        """'add members' requires two arguments. Both 'project_filter' and 'group_filter' can be either the full
        project/group name or a partial name which matches some projects/groups. 
        You can add SA emails from multiple projects to multiple groups if you wish."""
        admin = self._svc(*ADMIN, group_token)
        projId_list = self._list.projects(project_filter, file_tag, token, prt)
        group_list = self._list.groups(group_filter, file_tag, group_token, prt)
        for group in group_list:
            for project in projId_list:
                batch = BatchJob(admin)
                sa_emails, _ = self._list.sas(project, True, file_tag, token, False)
                if len(sa_emails) == 0:
                    print("No service accounts found in", project)
                    continue
                group_members = self._list.members(group, file_tag, group_token, prt)
                add_sas = [i for i in sa_emails if i not in group_members]
                if len(add_sas) == 0:
                    continue
                print(
                    f"Adding {str(len(add_sas))} SAs to group: {group} from project: {project}"
                )
                retries = retry
                while add_sas and retries:
                    retries -= 1
                    for email in add_sas:
                        batch.add(
                            admin.members().insert(
                                groupKey=group, body={"email": email, "role": "MEMBER"}
                            )
                        )
                    batch.execute()
                    sleep(sleep_time)
                    group_members = self._list.members(
                        group_filter, file_tag, group_token, prt
                    )
                    add_sas = [i for i in sa_emails if i not in group_members]
                print(
                    f"{len(sa_emails)-len(add_sas)} SA emails from {project} are in {group}. {len(add_sas)} failed."
                )

    def user(self, user, *td_id, role="organizer"):
        """Add user (typically group name) to shared/team drive(s). Usage: 'safire add mygroup@domain.com td_filter'"""
        for td_filt in td_id:
            drives = self._list.drives(td_filt, "", token, False)
            for td in drives:
                if td_filt in td['name']:
                    print(f"Adding {user} to {td['name']} {td['id']}")
                    drive = self._svc(*DRIVE, token)
                    body = {"type": "user", "role": role, "emailAddress": user}
                    drive.permissions().create(body=body, fileId=td['id'], supportsAllDrives=True, fields="id").execute()
                else:
                    pass

"""###### Remove"""

class Remove(Help):
    """Delete sas, jsons/keys, drives and group members. Note: 'remove' and 'delete' are equivalent commands"""

    def __init__(self):
        super(Remove, self).__init__()
        self._list = List()

    def sas(
        self,
        project_filter,
        token=token,
        exact=False,
        file_tag="",
        prt=False,
        retry=retry,
    ):
        """Usage: 'safire remove sas filter' where filter is a string to match the projects from which 
        you want to delete service accounts. To remove all SAs for all projects use "" as your filter"""
        projId_list = self._list.projects(project_filter, file_tag, token, prt)
        iam = self._svc(*IAM, token)
        for project in projId_list:
            sas, _ = self._list.sas(project, True, file_tag, token, False)
            retries = retry
            while sas and retries:
                retries -= 1
                if len(sas) == 0:
                    print(f"0 service accounts in {project}. Moving to next project")
                    continue
                batch = BatchJob(iam)
                print(f"Removing {len(sas)} service accounts from {project}")
                for i in sas:
                    name = f"projects/{project}/serviceAccounts/{i}"
                    batch.add(iam.projects().serviceAccounts().delete(name=name))
                batch.execute()
                sleep(sleep_time)
                sas, _ = self._list.sas(project, True, file_tag, token, False)

    def jsons(self, filter="", sa_path=sa_path):
        """Remove json keys from svcaccts path"""
        _, _, files = next(os.walk(sa_path))
        if not files:
            return f"No json keys found in {sa_path}, Nothing to delete"
        delete_files = [i for i in sorted(files) if ".json" and str(filter) in i]
        # print(f"Files to be deleted:\n", *delete_files, sep="\n")
        # yes_no = input("Confirm you want to delete all of these files. y/n: ")
        # if yes_no.lower() in ["y", "yes"]:
        for file in delete_files:
            print(f"Removing {file}")
            os.remove(f"{sa_path}/{file}")
        # else:
        #     print("Deletion of json files aborted")

    def members(
        self,
        group_filter,
        retry=retry,
        batch_size=50,
        file_tag="",
        group_token=group_token,
        prt=False,
    ):
        """Remove members from groups. Match 'filter'"""
        admin = self._svc(*ADMIN, group_token)
        batch = BatchJob(admin)
        group_list = self._list.groups(group_filter, file_tag, group_token, prt)
        group_members = []
        for group in group_list:
            group_members = self._list.members(group, file_tag, group_token, prt)
            retries = retry^2
            while len(group_members) > 1 and retries:
                print(
                    f"Removing {str(len(group_members))} SAs from group: {group}. Batch = {batch_size}"
                )
                retries -= 1
                for member in group_members[:batch_size]:
                    batch.add(admin.members().delete(groupKey=group, memberKey=member))
                batch.execute()
                batch = BatchJob(admin)
                sleep(sleep_time/4)
                group_members = self._list.members(group, file_tag, group_token, prt)
            print(
                f"{len(group_members)} members remaining in {group} (excluding OWNER)"
            )

    def drive(self, teamDriveId, token=token):
        """Delete a team/shared drive. Usage: 'safire add teamdrive unique ID'. USE CAREFULLY!
        Does not work with non-empty drives."""
        drvsvc = self._svc(*DRIVE, token)
        drives = self._list.drives(teamDriveId, "", token, False)
        print("Drives to be removed:", *drives, sep="\n")
        yes_no = input("Confirm you want to delete all of these drives. y/n: ")
        if not yes_no.lower() in ["y", "yes"]:
            return "Deletion of drives aborted"
        for drive in drives:
            print(f"Removing drive: {drive['name']} with id: {drive['id']}")
            drvsvc.drives().delete(driveId=str(drive["id"])).execute()

    def drives(self, input_file, token=token):
        """Delete team/shared drives. Usage: 'safire add teamdrive some_filename' with unique IDs. USE CAREFULLY"""
        td_list = open(input_file, "r")
        for teamDriveId in td_list:
            print(f"Deleting {teamDriveId}")
            self.drive(teamDriveId.rstrip())

    def user(self, user, *td_id, role="organizer", token=token):
        """Remove user (typically group name) from shared/team drive(s). Usage: 'safire remove mygroup@domain.com td_filter'"""
        drives = self._list.drives("", "", token, False)
        for td_filt in td_id:
            for td in drives:
                if td_filt in td['name']:
                    drvsvc = self._svc(*DRIVE, token)
                    user_info = []
                    resp = {'nextPageToken':None}
                    while 'nextPageToken' in resp:
                        resp = drvsvc.permissions().list(fileId=td['id'],pageSize=100,fields='nextPageToken,permissions(id,emailAddress,role)',supportsAllDrives=True,pageToken=resp['nextPageToken']).execute()
                        user_info += resp['permissions']
                    for i in user_info:
                        if user in i['emailAddress']:
                            print(f"Removing {user} with id {i['id']} from {td['name']} {td['id']}")
                            drvsvc.permissions().delete(permissionId=i['id'], fileId=td['id'], supportsAllDrives=True, fields="id").execute()

"""###### Rename"""

class Rename:
    """Rename json/key files to their email prefix, email numeric (omit prefix), uniqId or in a sequence.
    Usage: 'safire rename jsons email'  [choice email, email_seq, uniq, seq]
    Renaming is repeatable. Can always delete and redownload keys if needed."""
    def jsons(self, rename_type, dir=f"{sa_path}/"):
        """Usage: 'safire rename jsons email'  [choice email, email_seq, uniq, seq]"""
        import json, os
        filelist = os.listdir(dir)
        print("\nOriginal filenames:", *sorted(filelist), sep="\n")
        new_name = []
        if rename_type == "seq":
            [
                os.rename(dir + file, dir + f"{i+1}.json")
                for i, file in enumerate(sorted(filelist))
            ]
        else:
            for file in sorted(filelist):
                try:
                    with open(dir + file) as f:
                        data = json.load(f)
                        if rename_type == "email":
                            new_name = data["client_email"].split("@")[0] + ".json"
                        if rename_type == "email_seq":
                            digits = list(
                                filter(
                                    lambda x: x.isdigit(),
                                    data["client_email"].split("@")[0],
                                )
                            )
                            new_name = "".join(digits) + ".json"
                        if rename_type == "uniq":
                            new_name = data["client_id"] + ".json"
                        if os.path.exists(dir + new_name):
                            continue
                        os.rename(dir + file, dir + new_name)
                except:
                    continue
        print("\nCurrent filenames:", *sorted(os.listdir(dir)), sep="\n")

"""Add some command line tests

## Commands

- Enter command(s) and select `Runtime, Run All` from the top menu in Colab (or use ctrl-f9 to run all). 
- See below for results of each command. All results are also exported to xlsx and csv files in the /data directory
"""

# Commented out IPython magic to ensure Python compatibility.
# %%time

#@markdown
if "safire_ip.py" == sys.argv[0] and __name__ == "__main__":
    fire.Fire(Commands)
else:
#     %%time
    command1 = "" #@param {type:"string"}
    if command1:
        sys.argv = [""] + list(command1.split(" "))
        print(f"safire command: {command1} \n\nOutput: \n")
        fire.Fire(Commands)

    command2 = "" #@param {type:"string"}
    if command2:
        sys.argv = [""] + list(command2.split(" "))
        print(f"safire command: {command2} \n\nOutput: \n")
        fire.Fire(Commands)

    command3 = "" #@param {type:"string"}
    if command3:
        sys.argv = [""] + list(command3.split(" "))
        print(f"safire command: {command3} \n\nOutput: \n")
        fire.Fire(Commands)

    command4 = "" #@param {type:"string"}
    if command4:
        sys.argv = [""] + list(command4.split(" "))
        print(f"safire command: {command4} \n\nOutput: \n")
        fire.Fire(Commands)

    command5 = "" #@param {type:"string"}
    if command5:
        sys.argv = [""] + list(command5.split(" "))
        print(f"safire command: {command5} \n\nOutput: \n")
        fire.Fire(Commands)



"""###### Vars and values"""

# Commented out IPython magic to ensure Python compatibility.
#@markdown
# %whos