Skip to content

Commit

Permalink
disable ssh key generation and fabric_rc generation
Browse files Browse the repository at this point in the history
  • Loading branch information
@kthare10
kthare10 committed Jun 27, 2024
1 parent 8505ee6 commit b858120
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 67 deletions.
2 changes: 1 addition & 1 deletion jupyternb/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
__VERSION__ = "1.3.6"
__VERSION__ = "1.3.7b0"
__version__ = __VERSION__
166 changes: 100 additions & 66 deletions jupyternb/jupyter_startup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,9 +107,28 @@ def __init__(self):
if self.config_json_location is None:
self.config_json_location = self.DEFAULT_FABRIC_CONFIG_JSON_LOCATION

def create_config_dir(self):
def create_config_dir(self, dir_only: bool = True):
"""
Create a configuration directory and optionally populate environment files.
This method creates a directory specified by `self.config_location` and optionally populates
environment files based on environment variables. It sets up necessary configurations for
Fabric operations including credentials, keys, logging settings, and SSH configurations.
:param dir_only: If True (default), only create the directory without populating environment files.
If False, populate environment files based on configured environment variables.
:type dir_only: bool
:raises Exception: If there is an error during directory creation or file writing.
"""
try:
# Create the configuration directory
os.mkdir(self.config_location)

# If dir_only is True, return after creating the directory
if dir_only:
return

# Define environment variables with default or placeholder values
environment_vars = {
self.FABRIC_CREDMGR_HOST: os.environ[self.FABRIC_CREDMGR_HOST],
self.FABRIC_ORCHESTRATOR_HOST: os.environ[self.FABRIC_ORCHESTRATOR_HOST],
Expand All @@ -123,29 +142,38 @@ def create_config_dir(self):
self.FABRIC_LOG_LEVEL: self.DEFAULT_FABRIC_LOG_LEVEL,
self.FABRIC_LOG_FILE: self.DEFAULT_FABRIC_LOG_FILE
}

# Prepare string to write to environment file
string_to_write = ""
for key, value in environment_vars.items():
if '<' in value and '>':
string_to_write += f"#export {key}={value}\n"
if '<' in value and '>' in value:
string_to_write += f"#export {key}={value}\n" # Placeholder value comment
else:
string_to_write += f"export {key}={value}\n"
string_to_write += f"export {key}={value}\n" # Actual value

# Write environment file atomically
with atomic_write(f'{self.config_location}/fabric_rc', overwrite=True) as f:
f.write(string_to_write)

string_to_write = f"UserKnownHostsFile /dev/null\n" \
f"StrictHostKeyChecking no\n" \
f"ServerAliveInterval 120 \n" \
f"Host bastion-?.fabric-testbed.net\n" \
f"User <Update Bastion User Name>\n" \
f"ForwardAgent yes\n" \
f"Hostname %h\n" \
f"IdentityFile {self.config_location}/{os.environ[self.FABRIC_BASTION_PRIVATE_KEY_NAME]}\n" \
f"IdentitiesOnly yes\n" \
f"Host * !bastion-?.fabric-testbed.net\n" \
f"ProxyJump <Update Bastion User Name>@{os.environ[self.FABRIC_BASTION_HOST]}:22\n"
# Prepare SSH configuration string
ssh_config_string = (
f"UserKnownHostsFile /dev/null\n"
f"StrictHostKeyChecking no\n"
f"ServerAliveInterval 120 \n"
f"Host bastion-?.fabric-testbed.net\n"
f"User <Update Bastion User Name>\n"
f"ForwardAgent yes\n"
f"Hostname %h\n"
f"IdentityFile {self.config_location}/{os.environ[self.FABRIC_BASTION_PRIVATE_KEY_NAME]}\n"
f"IdentitiesOnly yes\n"
f"Host * !bastion-?.fabric-testbed.net\n"
f"ProxyJump <Update Bastion User Name>@{os.environ[self.FABRIC_BASTION_HOST]}:22\n"
)

# Write SSH configuration file atomically
with atomic_write(f'{self.config_location}/ssh_config', overwrite=True) as f:
f.write(string_to_write)
f.write(ssh_config_string)

except Exception as e:
print("Failed to create config directory and default environment file")
print("Exception: " + str(e))
Expand Down Expand Up @@ -269,36 +297,40 @@ def update_permissions(self):
"""
if not os.path.exists(f"{self.config_location}/fabric_rc"):
return
# Open the file in read mode
with open(f"{self.config_location}/fabric_rc", 'r') as file:
# Read the content of the file
content = file.read()

# Split the content into lines
lines = content.splitlines()

# Process each line to extract the environment variable and its value
for line in lines:
# Ignore empty lines or lines starting with '#' (comments)
if line.strip() == '' or line.strip().startswith('#'):
continue

# Split each line into variable and value using the first occurrence of '='
variable, value = line.split('=', 1)
variable = variable.replace("export", "")

# Remove leading/trailing whitespaces from the variable and value
variable = variable.strip()
value = value.strip()

if variable in [self.FABRIC_BASTION_KEY_LOCATION, self.FABRIC_SLICE_PRIVATE_KEY_FILE]:
if os.path.exists(value):
os.chmod(value, 0o600)
elif variable in [self.FABRIC_SLICE_PUBLIC_KEY_FILE]:
if os.path.exists(value):
os.chmod(value, 0o644)

def initialize(self):
try:
# Open the file in read mode
with open(f"{self.config_location}/fabric_rc", 'r') as file:
# Read the content of the file
content = file.read()

# Split the content into lines
lines = content.splitlines()

# Process each line to extract the environment variable and its value
for line in lines:
# Ignore empty lines or lines starting with '#' (comments)
if line.strip() == '' or line.strip().startswith('#'):
continue

# Split each line into variable and value using the first occurrence of '='
variable, value = line.split('=', 1)
variable = variable.replace("export", "")

# Remove leading/trailing whitespaces from the variable and value
variable = variable.strip()
value = value.strip()

if variable in [self.FABRIC_BASTION_KEY_LOCATION, self.FABRIC_SLICE_PRIVATE_KEY_FILE]:
if os.path.exists(value):
os.chmod(value, 0o600)
elif variable in [self.FABRIC_SLICE_PUBLIC_KEY_FILE]:
if os.path.exists(value):
os.chmod(value, 0o644)
except Exception as e:
print(f"Exception occurred while updating permissions: {e}")
traceback.print_exc()

def initialize(self, sliver_keys: bool = False):
"""
Initialize Jupyter Notebook Container
"""
Expand All @@ -310,7 +342,7 @@ def initialize(self):

if not os.path.exists(self.config_location):
print("Creating config directory and all files")
self.create_config_dir()
self.create_config_dir(dir_only=True)

if not os.path.exists(self.requirements_location):
print("Creating default requirements.txt")
Expand All @@ -324,28 +356,30 @@ def initialize(self):
self.download_notebooks()

# Create SSH Keys
ssh_key = FABRICSSHKey.generate(comment="fabric@localhost", algorithm="rsa")
with atomic_write(f'{self.DEFAULT_PRIVATE_SSH_KEY}', overwrite=True) as f:
f.write(ssh_key.private_key)
with atomic_write(f'{self.DEFAULT_PUBLIC_SSH_KEY}', overwrite=True) as f:
f.write(f'{ssh_key.name} {ssh_key.public_key} {ssh_key.comment}')

# Default key in config directory
default_slice_priv_key_config = f'{self.config_location}/{os.environ[self.FABRIC_SLICE_PRIVATE_KEY_NAME]}'
default_slice_pub_key_config = f'{self.config_location}/{os.environ[self.FABRIC_SLICE_PUBLIC_KEY_NAME]}'

if not os.path.exists(default_slice_priv_key_config):
with atomic_write(default_slice_priv_key_config, overwrite=True) as f:
# This is now made optional as JH notebook takes care of key creation
if sliver_keys:
ssh_key = FABRICSSHKey.generate(comment="fabric@localhost", algorithm="rsa")
with atomic_write(f'{self.DEFAULT_PRIVATE_SSH_KEY}', overwrite=True) as f:
f.write(ssh_key.private_key)

if not os.path.exists(default_slice_pub_key_config):
with atomic_write(default_slice_pub_key_config, overwrite=True) as f:
with atomic_write(f'{self.DEFAULT_PUBLIC_SSH_KEY}', overwrite=True) as f:
f.write(f'{ssh_key.name} {ssh_key.public_key} {ssh_key.comment}')

os.chmod(self.DEFAULT_PRIVATE_SSH_KEY, 0o600)
os.chmod(self.DEFAULT_PUBLIC_SSH_KEY, 0o644)
os.chmod(default_slice_priv_key_config, 0o600)
os.chmod(default_slice_pub_key_config, 0o644)
# Default key in config directory
default_slice_priv_key_config = f'{self.config_location}/{os.environ[self.FABRIC_SLICE_PRIVATE_KEY_NAME]}'
default_slice_pub_key_config = f'{self.config_location}/{os.environ[self.FABRIC_SLICE_PUBLIC_KEY_NAME]}'

if not os.path.exists(default_slice_priv_key_config):
with atomic_write(default_slice_priv_key_config, overwrite=True) as f:
f.write(ssh_key.private_key)

if not os.path.exists(default_slice_pub_key_config):
with atomic_write(default_slice_pub_key_config, overwrite=True) as f:
f.write(f'{ssh_key.name} {ssh_key.public_key} {ssh_key.comment}')

os.chmod(self.DEFAULT_PRIVATE_SSH_KEY, 0o600)
os.chmod(self.DEFAULT_PUBLIC_SSH_KEY, 0o644)
os.chmod(default_slice_priv_key_config, 0o600)
os.chmod(default_slice_pub_key_config, 0o644)

self.update_permissions()
self.custom_install_packages()
Expand Down

0 comments on commit b858120

Please sign in to comment.