Tweaked #58 slightly. Execution is done with an attempt to retain file line numbers when executing. It also consolidates behavior of files and remote documents in a 'similar' manner.

3 files changed, 61 insertions, 49 deletions

diff --git a/archinstall/__main__.py b/archinstall/__main__.py
index 4575a7e1..e91e050e 100644
--- a/archinstall/__main__.py
+++ b/archinstall/__main__.py
@@ -1,48 +1,10 @@
-from urllib.parse import urlparse
 import archinstall
 import sys
 import os
-import glob
-import traceback
-import urllib.request
 
 # TODO: Learn the dark arts of argparse...
 #	   (I summon thee dark spawn of cPython)
 
-
-class ProfileNotFound(BaseException):
-	pass
-
-
-def find_examples():
-	"""
-	Used to locate the examples, bundled with the module or executable.
-
-	:return: {'guided.py' : './examples/guided.py', '<profile #2>' : '<path #2>'}
-	:rtype: dict
-	"""
-	cwd = os.path.abspath(f'{os.path.dirname(__file__)}')
-	examples = f"{cwd}/examples"
-
-	return {os.path.basename(path): path for path in glob.glob(f'{examples}/*.py')}
-
-
-def find(url):
-	parsed_url = urlparse(url)
-	if not parsed_url.scheme:
-		examples = find_examples()
-		if f"{url}.py" in examples:
-			return open(examples[f"{url}.py"]).read()
-		try:
-			return open(url, 'r').read()
-		except FileNotFoundError:
-			return ProfileNotFound(f"File {url} does not exist")
-	elif parsed_url.scheme in ('https', 'http'):
-		return urllib.request.urlopen(url).read().decode('utf-8')
-	else:
-		return ProfileNotFound(f"Cannot handle scheme {parsed_url.scheme}")
-
-
 def run_as_a_module():
 	"""
 	Since we're running this as a 'python -m archinstall' module OR
@@ -54,20 +16,13 @@ def run_as_a_module():
 		sys.argv.append('guided')
 
 	try:
-		profile = find(sys.argv[1])
-	except ProfileNotFound as err:
+		script = archinstall.find_installation_script(sys.argv[1])
+	except archinstall.ProfileNotFound as err:
 		print(f"Couldn't find file: {err}")
 		sys.exit(1)
 
 	os.chdir(os.path.abspath(os.path.dirname(__file__)))
-
-	try:
-		exec(profile)  # Is this is very safe?
-	except Exception as err:
-		print(f"Error in profile {sys.argv[1]}: {err}")
-		traceback.print_exc(file=sys.stdout)
-		sys.exit(1)  # Should prompt for another profile path instead
-
+	script.execute()
 		
 if __name__ == '__main__':
 	run_as_a_module()
diff --git a/archinstall/lib/exceptions.py b/archinstall/lib/exceptions.py
index 68ba8750..2a1cae14 100644
--- a/archinstall/lib/exceptions.py
+++ b/archinstall/lib/exceptions.py
@@ -5,4 +5,6 @@ class DiskError(BaseException):
 class ProfileError(BaseException):
 	pass
 class SysCallError(BaseException):
+	pass
+class ProfileNotFound(BaseException):
 	pass
\ No newline at end of file
diff --git a/archinstall/lib/profiles.py b/archinstall/lib/profiles.py
index 8a39e95e..015d41b7 100644
--- a/archinstall/lib/profiles.py
+++ b/archinstall/lib/profiles.py
@@ -1,5 +1,5 @@
 import os, urllib.request, urllib.parse, ssl, json, re
-import importlib.util, sys
+import importlib.util, sys, glob, hashlib
 from collections import OrderedDict
 from .general import multisplit, sys_command, log
 from .exceptions import *
@@ -42,6 +42,36 @@ def list_profiles(base='./profiles/', filter_irrelevant_macs=True):
 		break
 	return cache
 
+def find_examples():
+	"""
+	Used to locate the examples, bundled with the module or executable.
+
+	:return: {'guided.py' : './examples/guided.py', '<profile #2>' : '<path #2>'}
+	:rtype: dict
+	"""
+	cwd = os.path.abspath(f'{os.path.dirname(__file__)}')
+	examples = f"{cwd}/examples"
+
+	return {os.path.basename(path): path for path in glob.glob(f'{examples}/*.py')}
+
+def find_installation_script(profile):
+	parsed_url = urllib.parse.urlparse(profile)
+	if not parsed_url.scheme:
+		examples = find_examples()
+		if f"{profile}.py" in examples:
+			with open(examples[f"{profile}.py"]) as file:
+				return Script(file.read(), filename=os.path.basename(profile)+".py")
+		try:
+			with open(profile, 'r') as file:
+				return Script(file.read(), filename=os.path.basename(profile))
+		except FileNotFoundError:
+			return ProfileNotFound(f"File {profile} does not exist")
+	elif parsed_url.scheme in ('https', 'http'):
+		return Script(urllib.request.urlopen(profile).read().decode('utf-8'), filename=os.path.basename(profile))
+	else:
+		return ProfileNotFound(f"Cannot handle scheme {parsed_url.scheme}")
+
+
 class Imported():
 	def __init__(self, spec, imported):
 		self.spec = spec
@@ -56,6 +86,31 @@ class Imported():
 		if len(args) >= 2 and args[1]:
 			raise args[1]
 
+
+class Script():
+	def __init__(self, content, filename=''):
+		self.content = content
+		self.filename = filename
+
+	@property
+	def path(self):
+		temp_file_path = f"/tmp/{self.filename}_{hashlib.md5(os.urandom(12)).hexdigest()}.py"
+
+		with open(temp_file_path, "w") as temp_file:
+			temp_file.write(self.content)
+
+		return temp_file_path
+
+	def execute(self):
+		spec = importlib.util.spec_from_file_location(
+			"tempscript",
+			self.path
+		)
+		imported_path = importlib.util.module_from_spec(spec)
+		spec.loader.exec_module(imported_path)
+		sys.modules["tempscript"] = imported_path
+
+
 class Profile():
 	def __init__(self, installer, path, args={}):
 		self._path = path