Skip to content

Commit ee6b5a8

Browse files
committed
introducing configuration file for ship_a_key.py
1 parent 8bfb667 commit ee6b5a8

File tree

4 files changed

+49
-31
lines changed

4 files changed

+49
-31
lines changed

client/container_preparation/entrypoint.sh

+5-4
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ docker_path="/var/run/docker.sock"
1010
parse_args() {
1111
while [[ "$#" -gt 0 ]]; do
1212
case "$1" in
13+
--config) config="$2"; shift 2 ;;
1314
-b|--base-oci-image) base_oci_image="$2"; shift 2 ;;
1415
-s|--sif-path) sif_path="$2"; shift 2 ;;
1516
-e|--encrypted) encrypted=true; shift ;;
@@ -26,7 +27,7 @@ parse_args() {
2627
done
2728

2829
# Check for required arguments
29-
if [ -z "$base_oci_image" ] || [ -z "$sif_path" ] || [ -z "$data_path" ] || [ -z "$data_path_at_rest" ] || ( [ -z "$users" ] && [ -z "$groups" ] ) || [ -z "$compute_nodes" ]; then
30+
if [ -z "$config" ] || [ -z "$base_oci_image" ] || [ -z "$sif_path" ] || [ -z "$data_path" ] || [ -z "$data_path_at_rest" ] || ( [ -z "$users" ] && [ -z "$groups" ] ) || [ -z "$compute_nodes" ]; then
3031
echo echo "Please provides options for both of these programs : "
3132
python3 ./prepare_container.py --help
3233
python3 ./utils/ship_a_key.py --help
@@ -110,13 +111,13 @@ else
110111

111112
if [ -z "$users" ]; then
112113
# If the user provided only groups
113-
python3 ./utils/ship_a_key.py --username "$username" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
114+
python3 ./utils/ship_a_key.py --config $config --username "$username" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
114115
elif [ -z "$groups" ] ; then
115116
# If the user provided only users
116-
python3 ./utils/ship_a_key.py --username "$username" -u "$users" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
117+
python3 ./utils/ship_a_key.py --config $config --username "$username" -u "$users" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
117118
else
118119
# If the user provided both
119-
python3 ./utils/ship_a_key.py --username "$username" -u "$users" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
120+
python3 ./utils/ship_a_key.py --config $config --username "$username" -u "$users" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
120121
fi
121122

122123
echo -e "${YELLOW}[LUMI-SD]${NC}${BLUE}[Container preparation]${NC} Key written to the vault"

client/data_preparation/entrypoint.sh

+5-4
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
parse_args() {
88
while [[ "$#" -gt 0 ]]; do
99
case "$1" in
10+
--config) config="$2"; shift 2 ;;
1011
-i|--input-data) input_data="$2"; shift 2 ;;
1112
-o|--output-data) output_data="$2"; shift 2 ;;
1213
--data-path) data_path="$2"; shift 2 ;;
@@ -21,7 +22,7 @@ parse_args() {
2122
done
2223

2324
# Check for required arguments
24-
if [ -z "$input_data" ] || [ -z "$output_data" ] || [ -z "$data_path" ] || [ -z "$data_path_at_rest" ] || [ -z "$username" ] || ( [ -z "$users" ] && [ -z "$groups" ] ) || [ -z "$compute_nodes" ]; then
25+
if [ -z "$config" ] || [ -z "$input_data" ] || [ -z "$output_data" ] || [ -z "$data_path" ] || [ -z "$data_path_at_rest" ] || [ -z "$username" ] || ( [ -z "$users" ] && [ -z "$groups" ] ) || [ -z "$compute_nodes" ]; then
2526
echo echo "Please provides options for both of these programs : "
2627
python3 ./prepare_data.py --help
2728
python3 ./utils/ship_a_key.py --help
@@ -96,13 +97,13 @@ echo -e "${YELLOW}[LUMI-SD]${NC}${BLUE}[Data preparation]${NC} Writing key to th
9697
# Handle different cases of user provided compute nodes / user / groups
9798
if [ -z "$users" ]; then
9899
# If the user provided only groups
99-
python3 ./utils/ship_a_key.py --username "$username" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
100+
python3 ./utils/ship_a_key.py --config $config --username "$username" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
100101
elif [ -z "$groups" ] ; then
101102
# If the user provided only users
102-
python3 ./utils/ship_a_key.py --username "$username" -u "$users" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
103+
python3 ./utils/ship_a_key.py --config $config --username "$username" -u "$users" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
103104
else
104105
# If the user provided both
105-
python3 ./utils/ship_a_key.py --username "$username" -u "$users" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
106+
python3 ./utils/ship_a_key.py --config $config --username "$username" -u "$users" -g "$groups" -c "$compute_nodes" --data-path "$data_path" --data-path-at-rest "$data_path_at_rest" -i "$spiffeID" || end_entrypoint "$spire_agent_pid" 1
106107
fi
107108

108109
echo -e "${YELLOW}[LUMI-SD]${NC}${BLUE}[Data preparation]${NC} Key written to the vault"

utils/ship_a_key.py

+35-21
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from hashlib import sha512
1212
from ssh_utils import ssh_connect, ssh_copy_file
1313

14+
from configparser import ConfigParser, NoSectionError, NoOptionError
15+
1416
# Provide client_id from cli$
1517
# Same for trust domain
1618
# Get image id and transform as for server
@@ -26,6 +28,11 @@ def parse_arguments() -> argparse.ArgumentParser:
2628
"""
2729
parser = argparse.ArgumentParser(description="CLI Options")
2830

31+
parser.add_argument(
32+
"--config",
33+
required=True,
34+
help="Path to the client configuration file",
35+
)
2936
parser.add_argument(
3037
"--users",
3138
"-u",
@@ -82,19 +89,6 @@ def parse_arguments() -> argparse.ArgumentParser:
8289
type=str,
8390
help="Path to write the dataset on the supercomputer storage default :",
8491
)
85-
parser.add_argument(
86-
"--sd-server-address",
87-
"-a",
88-
type=str,
89-
help="Server address",
90-
)
91-
parser.add_argument(
92-
"--sd-server-port",
93-
"-ap",
94-
type=int,
95-
default=10080,
96-
help="SD API server port (default: 10080)",
97-
)
9892
parser.add_argument(
9993
"--username",
10094
required=True,
@@ -103,6 +97,25 @@ def parse_arguments() -> argparse.ArgumentParser:
10397

10498
return parser.parse_args()
10599

100+
# Parse configuration file
101+
def parse_configuration(path : str):
102+
config = ConfigParser()
103+
config.read(path)
104+
105+
if not 'hpcs-server' in config:
106+
raise NoSectionError("hpcs-server section missing in configuration file, aborting")
107+
108+
if not 'vault' in config:
109+
raise NoSectionError("vault section missing in configuration file, aborting")
110+
111+
if not 'url' in config['hpcs-server']:
112+
raise NoOptionError("'hpcs-server' section is incomplete in configuration file, aborting")
113+
114+
if not 'url' in config['vault']:
115+
raise NoOptionError("'vault' section is incomplete in configuration file, aborting")
116+
117+
return config
118+
106119

107120
def validate_options(options: argparse.ArgumentParser):
108121
"""Check for the cli-provided options
@@ -194,7 +207,7 @@ def validate_options(options: argparse.ArgumentParser):
194207

195208

196209
def create_authorized_workloads(
197-
SVID: JwtSvid, secret, server, port, users, groups, compute_nodes
210+
SVID: JwtSvid, secret, url, users, groups, compute_nodes
198211
):
199212
"""Create workloads that are authorized to access to a secret
200213
@@ -212,7 +225,7 @@ def create_authorized_workloads(
212225
"""
213226

214227
# Prepare request
215-
url = f"http://{server}:{port}/api/client/create-workloads"
228+
url = f"{url}/api/client/create-workloads"
216229
payload = {
217230
"jwt": SVID.token,
218231
"secret": secret,
@@ -248,7 +261,9 @@ def create_authorized_workloads(
248261

249262
if __name__ == "__main__":
250263
# Parse arguments from CLI
251-
options = parse_arguments()
264+
options = parse_arguments()
265+
# Parse configuration file
266+
configuration = parse_configuration(options.config)
252267

253268
# Validate / Parse them
254269
(
@@ -277,15 +292,14 @@ def create_authorized_workloads(
277292
users_spiffeID, client_id, secrets_path, user_role = create_authorized_workloads(
278293
SVID,
279294
secret_name,
280-
options.sd_server_address,
281-
options.sd_server_port,
295+
configuration["hpcs-server"]["url"],
282296
users,
283297
groups,
284298
compute_nodes,
285299
)
286300

287301
# Login to the vault using client's certificate
288-
hvac_client = vault_login(SVID, f"client_{client_id}")
302+
hvac_client = vault_login(configuration["vault"]["url"], SVID, f"client_{client_id}")
289303

290304
# Prepare secret
291305
secret = {}
@@ -329,11 +343,11 @@ def create_authorized_workloads(
329343
ssh_copy_file(
330344
ssh_client,
331345
"/tmp/dataset_info.yaml",
332-
f"{options.data_path_at_rest}{secret_name}.info.yaml",
346+
f"{options.data_path_at_rest}/{secret_name}.info.yaml",
333347
)
334348

335349
print(
336-
f"Data and info file were shipped to te supercomputer. Infos about the dataset are available at {options.data_path_at_rest}/{secret_name}.info.yaml"
350+
f"Data and info file were shipped to te supercomputer. Info about the dataset are available at {options.data_path_at_rest}/{secret_name}.info.yaml"
337351
)
338352

339353
ssh_client.close()

utils/ssh_utils.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from paramiko.client import SSHClient
2-
from paramiko import SSHException, AutoAddPolicy
2+
from paramiko import SSHException, AutoAddPolicy, RSAKey
33
from scp import SCPClient
44

55
# Hostname and port configuration
@@ -30,11 +30,13 @@ def ssh_connect(username: str) -> SSHClient:
3030

3131
# Probably running in a container
3232
except SSHException:
33+
pkey=RSAKey.from_private_key_file("/tmp/.ssh/id_rsa")
3334
client.connect(
3435
host,
3536
port,
3637
username=username,
37-
key_filename="/tmp/.ssh/id_rsa",
38+
pkey=pkey,
39+
look_for_keys=False,
3840
auth_timeout=30,
3941
timeout=30,
4042
)

0 commit comments

Comments
 (0)