#!/usr/bin/python3
#
###############################################################################
##   Publish ISO Image To AOS Image Store or ESX Datastore
##   Filename: publish_iso_image_to_cluster.py
##   Script Version: 1.0.0
##	URL: https://www.nutanix.dev/code_samples/publish-iso-image-to-aos-image-store-or-esx-datastore
##############################################################################
#.NOTES
#	1. The AOS image store exists ONLY if the Nutanix cluster is running AHV as a Hypervisor.  Nutanix clusters running ESXi Hypervisors use ESX datastores, so the AOS image store does not exist.
#.PREREQUISITES
#	1. Two Linux Servers, one where this script runs, one that is running as a Web Server to publish and stage the ISO image for AOS image store URL import.  Conceivably, these two servers can be consolidated into one.
#	2. Nutanix Prism Element credentials (not Prism Central)
#	3. ESXi root (or other SSH enabled user) credentials, 
#	4. Staging Web Server user credentials.
#	5. Appropriate file level permissions to all filesystems above for handling file copy processes under provided user credentials.
#	6. Python3 installed with the below defined (import) modules.  Use an updated PIP3 to install missing modules (this is required for the latest cryptography)
#.SYNOPSIS
#	This script will take an ISO located local to the server executing the script and either a) copy the ISO to the published directory on the staging Web Server then use the published URL
#	to envoke a Prism REST call to import the ISO or b) copy the ISO to the defined ESXi datastore location.
#.CONFIGURATION
#	1. Use the Python3 console to encrypt credentials to be used in defining variables below.
#		[user@host ~:] python3
#		>>> import base64
#		>>> print(base64.b64encode("mypass".encode("utf-8"))) <-- Define the results from this in variable. Copy everything inside the ' and ', nothing outside.
#		>>> print(base64.b64decode("bXlwYXNz").decode("utf-8")) <-- Validate the above encrypted password using this (if needed)
#		>>> quit()
#	2. Configure environment specific variables below.
#.EXECUTION
# usage: publish_iso_image_to_cluster.py [-h] -t [aos|esx] -i /path/to/isofile.iso [-a Alternate Name] [-d Image Description]
#		This script will upload ISO images to ESX Datastores or an AOS Image Store.
#		optional arguments:
#				-h, --help --> Show this help message and exit.
#				-t [aos|esx], --type [aos|esx] --> Choose the platform to target.
#				-i /path/to/isofile.iso, --isofile /path/to/isofile.iso --> Specify the path to the ISO file.
#				-a Alternate Name, --altname Alternate Name --> Specify an alternate name for the ISO.
#				-d Image Description, --description Image Description --> Specify a description comment for the ISO. Only works with AOS.
#.EXAMPLE
#	[user@localhost.localdomain ~]$ python3 ./publish_iso_image_to_cluster.py -t aos -i /home/imageuser/Test1.iso -a "Testing" -d "Blah"
#	
#	Target Platform: AOS
#	ISO Name: Testing
#	ISO Description: Blah
#	ISO File: /home/imageuser/Test1.iso
#	ISO Hash: 1b40441f8555d1d56dfce928b8d72791
#	
#	Uploading Testing.iso to Staging Web Server: xx.xxx.xx.123 [100.00%]
#	Remote ISO Hash: 1b40441f8555d1d56dfce928b8d72791
#	
#	Local and Remote file hash matches...
#	Image upload to Staging Web Server complete!
#	
#	Publishing Testing.iso to AOS Image Store (xx.xxx.xx.456)
#	Publish successful! [201]
#.DISCLAIMER
#	This code is intended as a standalone example.  Subject to licensing restrictions defined on nutanix.dev, this can be downloaded, copied and/or modified in any way you see fit.
#	Please be aware that all public code samples provided by Nutanix are unofficial in nature, are provided as examples only, are unsupported and will need to be heavily scrutinized and potentially modified before they can be used in a production environment.  All such code samples are provided on an as-is basis, and Nutanix expressly disclaims all warranties, express or implied.
# 
#	All code samples are © Nutanix, Inc., and are provided as-is under the MIT license. (https://opensource.org/licenses/MIT)
##############################################################################
# SET VARIABLES BELOW
##############################################################################
MY_MD5SUM_PATH = '/usr/bin/md5sum'
# AOS CLUSTER (PRISM VIP)
MY_AOS_CLUSTER_VIP = "xx.xxx.xx.xxx"
MY_AOS_CLUSTER_USER = "prism_user_account_with_admin_permissions"
MY_AOS_CLUSTER_PASS = "bXlwYXNz"
MY_AOS_STORAGE_CONTAINER_ID = "219146243"
# ESXI HOST SSH (not vCenter)
MY_ESX_HOST_IP = "xx.xxx.xx.xxx"
MY_ESX_HOST_USER = "esxi_host_user_ssh_and_filesystem_write_permissions"
MY_ESX_CLUSTER_PASS = "bXlwYXNz"
MY_ESX_DATASTORE_PATH = "/vmfs/volumes/container/images_directory"
# Web Server SSH
MY_WEBSERVER_HOST_IP = "xx.xxx.xx.xxx"
MY_WEBSERVER_ISO_PATH = "/var/www/html/isoimages"
MY_WEBSERVER_USER = "webserver_username"
MY_WEBSERVER_PASS = "bXlwYXNz"
##############################################################################
#////////////////////////////////////////////////////////////////////////////////////////////////
# CHANGE NOTHING BELOW HERE!
#////////////////////////////////////////////////////////////////////////////////////////////////
##############################################################################
import argparse
import base64
import json
import os
import paramiko
import re
import requests
import subprocess
import sys
import urllib3
from paramiko import SSHClient
from pathlib import Path
from scp import SCPClient
from requests.auth import HTTPBasicAuth
parser = argparse.ArgumentParser(
    description='This script will upload ISO images to ESX Datastores or an AOS Image Store.',
	add_help=False
)

parser.add_argument('-h', '--help', action='help', default=argparse.SUPPRESS, help='Show this help message and exit.')
parser.add_argument('-t', '--type', metavar='[aos|esx]', required=True, choices={'aos','esx'}, type=str.lower, help='Choose the platform to target.')
parser.add_argument('-i', '--isofile', metavar='/path/to/isofile.iso', required=True, help='Specify the path to the ISO file.')
parser.add_argument('-a', '--altname', metavar='Alternate Name', required=False, help='Specify an alternate name for the ISO.')
parser.add_argument('-d', '--description', metavar='Image Description', required=False, help='Specify a description comment for the ISO. Only works with AOS.')

args = parser.parse_args()

MY_ISO_FILE = str(args.isofile)
MY_PATH = Path(MY_ISO_FILE)
MY_ISO_NAME_REGEX_PATTERN = "^[a-z0-9]+([-_\s]{1}[a-z0-9]+)*$"
MY_ISO_DESCRIPTION = ""
MY_ISO_FILENAME = MY_PATH.name

if args.altname:
	if bool(re.match(MY_ISO_NAME_REGEX_PATTERN, str(args.altname), re.IGNORECASE)):
		MY_ISO_NAME = str(args.altname)
		MY_ISO_FILENAME = MY_ISO_NAME + MY_PATH.suffix
	else:
		MY_ISO_NAME = MY_PATH.stem
		MY_ISO_FILENAME = MY_ISO_NAME + MY_PATH.suffix
else:
	MY_ISO_NAME = MY_PATH.stem

if args.description and str(args.type) == 'aos':
	if bool(re.match(MY_ISO_DESCRIPTION, str(args.description), re.IGNORECASE)):
		MY_ISO_DESCRIPTION = str(args.description)
	else:
		MY_ISO_DESCRIPTION = str('')
else:
	MY_ISO_DESCRIPTION = str('')

def upload():
	try:
		if args.type == 'aos':
			ssh = SSHClient()
			ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
			ssh.connect(MY_WEBSERVER_HOST_IP, '22', MY_WEBSERVER_USER, base64.b64decode(MY_WEBSERVER_PASS).decode("utf-8"))
			def progress(filename, size, sent):
				sys.stdout.write("Uploading " + MY_ISO_FILENAME + ' to Staging Web Server: ' + MY_WEBSERVER_HOST_IP + " [%.2f%%]  \r" % (float(sent)/float(size)*100) )
			scp = SCPClient(ssh.get_transport(), progress = progress)
			scp.put(MY_ISO_FILE,MY_WEBSERVER_ISO_PATH + '/' + MY_ISO_FILENAME)
			scp.close()
			print('\r')
			command = "md5sum " + MY_WEBSERVER_ISO_PATH + '/' + MY_ISO_FILENAME + ' | awk \'{print $1}\''
			stdin, stdout, stderr = ssh.exec_command(command)
			for line in iter(stdout.readline, ""):
				MY_REMOTE_ISO_HASH = line
			print('Remote ISO Hash: ' + MY_REMOTE_ISO_HASH)
			return MY_REMOTE_ISO_HASH
		elif args.type == 'esx':
			ssh = SSHClient()
			ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
			ssh.connect(MY_ESX_HOST_IP, '22', MY_ESX_HOST_USER, base64.b64decode(MY_ESX_CLUSTER_PASS).decode("utf-8"))
			def progress(filename, size, sent):
				sys.stdout.write("Uploading " + MY_ISO_FILENAME + ' to ESX Datastore: ' + MY_ESX_HOST_IP + " [%.2f%%]  \r" % (float(sent)/float(size)*100) )
			scp = SCPClient(ssh.get_transport(), progress = progress)
			scp.put(MY_ISO_FILE,MY_ESX_DATASTORE_PATH + '/' + MY_ISO_FILENAME)
			scp.close()
			print('\r')
			command = "md5sum " + MY_ESX_DATASTORE_PATH + '/' + MY_ISO_FILENAME + ' | awk \'{print $1}\''
			stdin, stdout, stderr = ssh.exec_command(command)
			for line in iter(stdout.readline, ""):
				MY_REMOTE_ISO_HASH = line
			print('Remote ISO Hash: ' + MY_REMOTE_ISO_HASH)
			return MY_REMOTE_ISO_HASH
	except: # catch *all* exceptions
		e = sys.exc_info()[0]
		print( "\nException: %s" % e )
		sys.exit()

def do_aos():
	MY_ISO_URI = 'http://' + MY_WEBSERVER_HOST_IP + '/' + MY_WEBSERVER_ISO_PATH.split('/')[-1] + '/' + MY_ISO_FILENAME

	# Disable insecure connection warnings please be advised and aware of the implications of doing this in a production environment!
	urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

	endpoint = 'https://' + MY_AOS_CLUSTER_VIP + ':9440/PrismGateway/services/rest/v2.0/images/'
	print('Publishing ' + MY_ISO_FILENAME + ' to AOS Image Store (' + MY_AOS_CLUSTER_VIP + ')\r')

	request_headers = {"Content-Type": "application/json", "charset": "utf-8"}
	request_payload = {
		"annotation": MY_ISO_DESCRIPTION,
		"image_import_spec": {
			"storage_container_id": MY_AOS_STORAGE_CONTAINER_ID,
			"url": MY_ISO_URI
		},
		"image_type": "ISO_IMAGE",
		"name": MY_ISO_NAME
	}

	try:
		results = requests.post(
			endpoint,
			data=json.dumps(request_payload),
			headers=request_headers,
			verify=False,
			auth=HTTPBasicAuth(MY_AOS_CLUSTER_USER, base64.b64decode(MY_AOS_CLUSTER_PASS).decode("utf-8")),
		)
		sc = str(results.status_code) 
		# check the results of the request
		if results.status_code == 200 or results.status_code == 201:
			print("Publish successful! [" + sc + "]\n")

	except: # catch *all* exceptions
		e = sys.exc_info()[0]
		print( "\nException: %s" % e )
		sys.exit()

if MY_ISO_FILE.lower().endswith('.iso'):
	if os.path.isfile(MY_ISO_FILE):
		ps = subprocess.Popen(MY_MD5SUM_PATH + " " + MY_ISO_FILE + " | awk '{print $1}'",shell=True,stdout=subprocess.PIPE,stderr=subprocess.STDOUT)
		MY_ISO_HASH = str(ps.communicate()[0],'utf-8')
		if args.type == 'aos':
			print('\nTarget Platform: AOS\nISO Name: ' + MY_ISO_NAME)
			if MY_ISO_DESCRIPTION:
				print('ISO Description: ' + MY_ISO_DESCRIPTION)
			print('ISO File: ' + MY_ISO_FILE + '\nISO Hash: ' + MY_ISO_HASH)
			MY_REMOTE_ISO_HASH =  upload()
			if MY_REMOTE_ISO_HASH == MY_ISO_HASH:
				print('Local and Remote file hash matches...\nImage upload to Staging Web Server complete!\n')
				do_aos()
			else:
				print('Abort: Local and Remote file hash do not match!\nImage upload to Staging Web Server failed!\n')
				sys.exit()
		elif args.type == 'esx':
			print('\nTarget Platform: ESX\nISO Name: ' + MY_ISO_NAME + '\nISO File: ' + MY_ISO_FILE + '\nISO Hash: ' + str(MY_ISO_HASH))
			MY_REMOTE_ISO_HASH =  upload()
			if MY_REMOTE_ISO_HASH == MY_ISO_HASH:
				print('Local and Remote file hash matches...\nImage upload to ESX Datastore complete!\n')
			else:
				print('Abort: Local and Remote file hash do not match!\nImage upload to ESX Datastore failed!\n')
	else:
		print('Error: file does not exist.\n')
else:
	print('Error: file must be an ISO file.\n')