#!/usr/local/bin/python3
# Tim Standing
# 22.10.04

import sys
import os
import subprocess
import pathlib
import time

kInvalidState = 0
kCheckInternetState = 1
kCheckUsernameAndPasswordState = 2
kCheckTargetIsBundleState = 3
kCompressTargetState = 4
kUploadTargetState = 5
kCheckProgressState = 6
kWaitTenSecondsState = 7
kSuccessState = 8
kFailureState = 9
kDoneState = 10

kNotarizationInProgress = 0
kNotarizationSucceeded = 3
kNotarizationFailed = 2

kLengthOfUUID_String = 36

#	====================================================================================
#
#	====================================================================================
def check_internet_connection():
	result = True
	
	print("----> checking internet connection")
	
	command = ["/sbin/ping", "-t", "2", "-q", "apple.com"]
	completedProcessObject = subprocess.run(command, capture_output=True, text=True)
	
	returnCode = completedProcessObject.returncode
	
	if returnCode != 0:
		result = False
		print ("####> Error: no internet connection available.")

	stdoutString = str(completedProcessObject.stdout)
	stderrString = str(completedProcessObject.stderr)
	
	return result

#	====================================================================================
#	Use notarytool to see if the username, password and team ID passed in are correct.
#	====================================================================================
def check_username_and_password(appleID, password, teamID):
	result = True

	print("----> checking username and password")
	
	if result:
		command = ["/usr/bin/xcrun", "notarytool", "history", "--apple-id", appleID,
			"--password", password, "--team-id", teamID]
		completedProcessObject = subprocess.run(command, capture_output=True, text=True)
	
		returnCode = completedProcessObject.returncode
			
		if returnCode != 0:
			print ("####> Error: apple id, application password or team ID invalid.")
			result = False
			
	return result

#	====================================================================================
#
#	====================================================================================
def is_valid_target_file(path):
	result = True
	
	result = len(path.suffixes) != 0
	
	if result:
		extension = path.suffixes[len(path.suffixes) - 1]
		result = extension == ".dmg" or extension == ".pkg"
		
	return result

#	====================================================================================
#
#	====================================================================================
def target_file_needs_compression(path):
	result = True
	
	result = len(path.suffixes) != 0
	
	if result:
		extension = path.suffixes[len(path.suffixes) - 1]
		result = extension == ".pkg"
	
	return result

#	====================================================================================
#
#	====================================================================================
def is_valid_target_bundle(path):
	result = True
	
	result = len(path.suffixes) > 0
	
	if result:
		extension = path.suffixes[len(path.suffixes) - 1]
		result = extension == ".app" or extension == ".framework" or extension == ".kext" or extension == ".plugin"
		
	return result

#	====================================================================================
#
#	====================================================================================
def compress_target(path):
	result = pathlib.Path(str(path.parent) + "/Archive_" + str(path.name) + ".zip")
	
	print("----> compress target")
	
	result.unlink(True)
	os.chdir(str(path.parent))
	command = ["/usr/bin/zip", "-y", "-r", str(result), str(path)]
	
	completedProcessObject = subprocess.run(command, capture_output=True, text=True)
		
	return result

#	====================================================================================
#
#	====================================================================================
def upload_target(appleID, password, teamID, path):
	result = ""
	print("----> upload target")
	
	command = ["/usr/bin/xcrun", "notarytool", "submit", "--apple-id", appleID, "--password", password, "--team-id", teamID, str(path)]
	completedProcessObject = subprocess.run(command, capture_output=True, text=True)
	
	returnCode = completedProcessObject.returncode
	stdoutString = str(completedProcessObject.stdout)
	stderrString = str(completedProcessObject.stderr)

	offset = stdoutString.find("id: ")
	offset = offset + len("id: ")
	result = stdoutString[offset : (offset + kLengthOfUUID_String)]

	if (len(result) > 0):
		print("----> upload ID =", result)
		
	return result

#	====================================================================================
#
#	====================================================================================
def check_notarization_progress(appleID, password, teamID, notarizationID):
	result = kNotarizationInProgress
	
	command = ["/usr/bin/xcrun", "notarytool", "info", notarizationID, "--apple-id", appleID, "--password", password, "--team-id", teamID]
	completedProcessObject = subprocess.run(command, capture_output=True, text=True)
	returnCode = completedProcessObject.returncode
	stdoutString = str(completedProcessObject.stdout)
	stderrString = str(completedProcessObject.stderr)
	
	if returnCode == 0:
		if (stdoutString.find("status: Accepted") >= 0):
			result = kNotarizationSucceeded
		elif (stdoutString.find("status: Invalid") >= 0):
			result = kNotarizationFailed
			print("####> Error: target rejected by Apple notarization server \n(most likely error is that target is not properly signed).\n\nnotarization server output =", stdoutString)
	
			command = ["/usr/bin/xcrun", "notarytool", "log", notarizationID, "--apple-id", appleID, "--password", password, "--team-id", teamID]
			completedProcessObject = subprocess.run(command, capture_output=True, text=True)
			returnCode = completedProcessObject.returncode
			stdoutString = str(completedProcessObject.stdout)
			print("notarization server log =", stdoutString)
			
		elif (stdoutString.find("status: In Progress") >= 0):
			result = kNotarizationInProgress
	else:
		result = kNotarizationFailed

	return result

#	====================================================================================
#	Staple the notarization
#	====================================================================================
def staple_notarization_to_target(path):
	command = ["/usr/bin/xcrun", "stapler", "staple", path]
	completedProcessObject = subprocess.run(command, capture_output=True, text=True)

#	====================================================================================
#
#	====================================================================================
def run_state_machine(appleID, password, teamID, pathname):
	result = 0
	done = False
	errorState = kInvalidState
	currentState = kInvalidState
	targetPath = pathlib.Path(pathname).resolve()
	uploadPath = targetPath
	totalSleepTime = 0
	
	while not done:
		#	Save the current state in case I encounter an error.
		errorState = currentState
		
		#	The switch statement which is the core of the state machine.
		if  currentState == kInvalidState:
			currentState = kCheckInternetState
			
		elif currentState == kCheckInternetState:
			if check_internet_connection():
				currentState = kCheckUsernameAndPasswordState
			else:
				currentState = kFailureState
		
		elif currentState == kCheckUsernameAndPasswordState:
			if check_username_and_password(appleID, password, teamID):
				currentState = kCheckTargetIsBundleState
			else:
				currentState = kFailureState

		elif currentState == kCheckTargetIsBundleState:
			if not targetPath.exists():
				currentState = kFailureState
				print ("####> Error: target file or bundle does not exist, at path: ", targetPath)
			elif targetPath.is_file():
				if is_valid_target_file(targetPath):
					if target_file_needs_compression(targetPath):
						currentState = kCompressTargetState
					else:
						currentState = kUploadTargetState
				else:
					currentState = kFailureState
					print ("####> Error: target file is not a .dmg or .pkg, at path: ", targetPath)
			elif targetPath.is_dir():
				if is_valid_target_bundle(targetPath):
					currentState = kCompressTargetState
				else:
					currentState = kFailureState
					print ("####> Error: target bundle is not a .app, .kext, .framework or .plugin, at path: ", targetPath)
			else:
				currentState = kFailureState
				print ("####> Error: invalid target, at path: ", targetPath)
				

		elif currentState == kCompressTargetState:
			uploadPath = compress_target(targetPath)
			if uploadPath.exists:
				currentState = kUploadTargetState
			else:
				currentState = kFailureState
				print ("####> Error: failure compressing target, at path: ", targetPath)

		elif currentState == kUploadTargetState:
			notarizationID = upload_target(appleID, password, teamID, uploadPath)
			if len(notarizationID) > 0:
				currentState = kCheckProgressState
			else:
				currentState = kFailureState
				print ("####> Error: failure uploading target to Apple's notarization server")
				
		elif currentState == kCheckProgressState:
			notarizationResult = check_notarization_progress(appleID, password, teamID, notarizationID)
			
			if notarizationResult == kNotarizationInProgress:
				currentState = kWaitTenSecondsState
			elif notarizationResult == kNotarizationSucceeded:
				currentState = kSuccessState
			elif notarizationResult == kNotarizationFailed:
				currentState = kFailureState

		elif currentState == kWaitTenSecondsState:
			totalSleepTime += 10
			print("----> waiting 10 seconds, total wait = ", totalSleepTime)
			time.sleep(10)
			currentState = kCheckProgressState

		elif currentState == kSuccessState:
			staple_notarization_to_target(targetPath)
			print ("----> notarization succeeded")
			currentState = kDoneState

		elif currentState == kFailureState:
			print ("----> notarization failed")
			result = 1
			currentState = kDoneState

		elif currentState == kDoneState:
			if uploadPath != targetPath:
				uploadPath.unlink(True)
				
			done = True
	
	return result
	
#	====================================================================================
#
#	====================================================================================
def print_help():
	print("\nnotarize_tool - tool for notarizing a .dmg, .pkg, or bundle with Apple's servers\nusage: notarize_tool apple_id password adc_team_id path\n\npassword is application specific password\nsee: https://support.apple.com/en-us/HT204397\n\nteam_id is your the Team ID found when you log into your developer account at https://developer.apple.com\n")
	
#	====================================================================================
#
#	====================================================================================
if __name__ == "__main__":
	result = 0
	
	argCount = len(sys.argv)
	
	if argCount != 5:
		print ("####> Error: incorrect number of parameters")
		print_help()
		result = 1
		
	if result == 0:
		appleID = sys.argv[1].lower()
		password = sys.argv[2]
		teamID = sys.argv[3]
		pathname = sys.argv[4]
	
		result = run_state_machine(appleID, password, teamID, pathname)
	
	sys.exit(result)