#!/usr/bin/python
import sys
import argparse
import json
import time

sys.path.insert(0, './classes/')
import cohesityAPI as cohesity
import serviceNowAPI as serviceNow
import sharePointAPI as sharePoint
import automationsAPI as dashboard

def GetArgs():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--server', '-s', type=str, action='store')
    parser.add_argument('--protectGroup', '-p', type=str, action='store')
    parser.add_argument('--help', '-h', action='store_true')
    parser.add_argument('--debug', '-d', action='store_true')
    return (parser.parse_args())

def PrintHelp():
    print("\nBasic Usage:")
    print("\npython3 serverRegistration.py -s hostname.nd.gov -p MyApp@physical")
    print("\t -s: Fully qualified domain name of the Server")
    print("\t -p: Protection Group the server will be in")
    print("\t -h Prints this help message")

args = GetArgs()

if not (args.server and args.protectGroup):
    PrintHelp()
    exit(1)

# Proceed with registration 
if (args.server and args.protectGroup):
    try:
        mdn = cohesity.API('itdmdndpc01.nd.gov')
        mdnToken = mdn.GetAuthToken()
        mdn.UpdateHeaders(mdnToken['accessToken'])

        newHost = args.server
        protectionGroup = args.protectGroup

        isPhysicalRegistered = bool()

        # Get the physical server objects in Cohesity
        physicalServers = mdn.GetFilteredRequest("/public/protectionSources/registrationInfo", "?environments=kPhysical")

        for nodes in physicalServers['rootNodes']:
            if nodes['rootNode']['name'] == newHost:
                print(newHost + " is already registerd as a physical server.")
                isPhysicalRegistered = bool("true")
                physicalSourceId = nodes['rootNode']['id']
                break

        if not isPhysicalRegistered:
            attempt = mdn.RegisterPhysical(newHost)
            print("Registering " + newHost + " as a Physical Server ")
            time.sleep(20)
            try:
                physicalSourceId = attempt['id']
            except:
                print("\nCould not reach " + newHost + " on port 50051. Please verify the agent is running and that all firewall rules are accounted for.")
                exit(1)
            
        # Get the list of protection groups & add the new one if necessary
        protectionJobs = mdn.GetFilteredRequest("/public/protectionJobs", "?environments[0]=kPhysical")
        uniqueJobs = {job['id'] : job for job in protectionJobs}.values()
        jobExists = bool()
        isProtected = bool()
        jobSources = []

        for job in uniqueJobs:
            # Source ID cannot be in more then on protection group for SQL
            if ('sourceIds' in job) and (physicalSourceId in job['sourceIds'] and ("DELETED" not in job['name']) ):
                print(newHost + " is already in " + job['name'] + " and cannot be in multiple groups.")
                isProtected = bool("true")
                jobExists = bool("true")
                break
            
            # The source ID was not found in any other SQL job, but a job name already exsits, so lets collect all sources in the job. 
            elif (job['name'] == protectionGroup) and ( physicalSourceId not in job['sourceIds']):
                jobExists = bool("true")
                break

        if jobExists == bool('true') and isProtected != bool('true'):
            print("Adding host " + newHost + " to protection group " + protectionGroup)

            # By default we will protect all local drives
            sourceParams = {
                "sourceId": physicalSourceId, 
                "physicalSpecialParameters": {
                    "filePaths": [
                        {
                            "backupFilePath":"$ALL_LOCAL_DRIVES",
                            "skipNestedVolumes":True
                        }
                    ],
                "usesSkipNestedVolumesVec":True
                }
            }

            job["sourceIds"].append(physicalSourceId)
            job["sourceSpecialParameters"].append(sourceParams)
            resp = mdn.UpdateProtectionJob(job)
            print(resp.content)

        if jobExists == bool():
            print("Creating protection job for " + protectionGroup + " and adding " + newHost)
            ################ TO DO ###############
            # Defaults from JSON dump of jobs objects
            resp = mdn.CreatePhysicalProtectionJob(physicalSourceId, protectionGroup)
            print(resp.content)
    
    except OSError as cohesityError:
        print('Cohesity Error: ' + cohesityError)

dashboard.send_automation({'AutomationName': 'Infra-Cohesity', 'Action': 'Provisioning', 'Platform': 'Python-registerServer.py', 'Units': 30})