Spaces:
Running
Running
File size: 5,631 Bytes
2a0bc63 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 |
"""
Facilities to manage the conversion between secure-connect-bundle (plus
additional info) and the "init string".
"""
import base64
import json
import os
import shutil
import tempfile
from typing import Any, Dict, Optional, Tuple
from zipfile import ZipFile
CONFIG_FILE_NAME = "config.json"
USED_CONFIG_KEYS_AND_DEFAULTS = {
"caCertLocation": "./ca.crt",
"keyLocation": "./key",
"certLocation": "./cert",
}
def encode_str(sng: str) -> str:
return base64.b64encode(sng.encode()).decode()
def decode_str(b64: str) -> str:
return base64.b64decode(b64).decode()
def _encode_from_string(s: str) -> str:
return encode_str(s)
def _encode_from_file(file_dir: str, file_title: str) -> str:
file_path = os.path.join(file_dir, file_title)
with open(file_path) as o_file:
return o_file.read()
def _clean_filename(fn: str) -> str:
if fn[:2] == "./":
return fn[2:]
else:
return fn
def infer_keyspace_from_bundle(bundle_path: Optional[str]) -> Optional[str]:
"""
Given a bundle zipfile path, try to peek at its config.json
and extract a default keyspace name on it.
Do not raise errors: rather, silently return None instead.
"""
if bundle_path:
try:
bundle_zip = ZipFile(bundle_path)
open_config = bundle_zip.open("config.json")
ascii_lines = [line.decode() for line in open_config.readlines()]
config = json.loads("".join(ascii_lines))
return str(config["keyspace"]) if "keyspace" in config else None
except Exception:
return None
else:
return None
def bundle_path_to_init_string(
secure_bundle: str,
keyspace: Optional[str] = None,
token: Optional[str] = None,
tempfile_basedir: Optional[str] = None,
) -> str:
"""
Utility function. Make the provided bundle zip into a "init string".
1. Unzip the provided bundle to a temp dir
2. Read and acquire (selected) contents of the zip
3. prepare the big json bundle description
4. make the latter into a big "init string"
5. delete the temporary dir
"""
with ZipFile(secure_bundle) as zipfile:
base_dir = tempfile_basedir if tempfile_basedir else tempfile.gettempdir()
temp_dir = tempfile.mkdtemp(dir=base_dir)
try:
# read config.json and find the filenames to get
zipfile.extract(CONFIG_FILE_NAME, path=temp_dir)
filenames = [CONFIG_FILE_NAME]
config = json.load(open(os.path.join(temp_dir, CONFIG_FILE_NAME)))
filenames += [
_clean_filename(config.get(cfg_k, default))
for cfg_k, default in USED_CONFIG_KEYS_AND_DEFAULTS.items()
]
for filename in filenames:
zipfile.extract(filename, path=temp_dir)
#
bundle_data = {
filename: _encode_from_file(temp_dir, filename)
for filename in filenames
}
#
bundle_description = {
"bundle_data": bundle_data,
"options": {
"version": "1",
"bundle_file_name": os.path.split(secure_bundle)[1],
"keyspace": keyspace,
"token": token,
},
}
#
init_string = encode_str(
json.dumps(
bundle_description,
separators=(",", ":"),
)
)
return init_string
finally:
shutil.rmtree(temp_dir)
def init_string_to_bundle_path_and_options(
init_string: str, target_dir: str
) -> Tuple[str, Dict[str, Any]]:
"""
Make an init string back into a usable secure-connect-bundle.
Return
(bundle_path, options_dict)
where options_dict is any additional info found in the init string
beyond the bundle contents.
"""
bundle_description = json.loads(decode_str(init_string))
options_dict = bundle_description.get("options", {})
version = options_dict.get("version")
if version == "1":
bundle_data = bundle_description["bundle_data"]
bundle_file_name = options_dict.get(
"bundle_file_name", "secure-connect-bundle.zip"
)
bundle_filepath = os.path.join(target_dir, bundle_file_name)
with ZipFile(bundle_filepath, "w") as bundle_zip_file:
for f_name, f_contents in bundle_data.items():
bundle_zip_file.writestr(f_name, data=f_contents)
return bundle_filepath, options_dict
else:
raise ValueError(f"Init string has unsupported or unknown version {version}.")
def create_init_string_utility() -> None:
# Utility command-line: converts the full bundle path (argument)
# to an "init string" and prints it as an export line
import sys
if sys.argv[1:] == [] or len(sys.argv) > 4:
cmd = sys.argv[0]
print(f"Usage: {cmd} bundle_path [keyspace [token]]")
else:
bundle_path = sys.argv[1]
keyspace = None if len(sys.argv) < 3 else sys.argv[2]
token = None if len(sys.argv) < 4 else sys.argv[3]
print(f'- Input bundle path: "{bundle_path}".')
print(f"- Keyspace={'Y' if keyspace else 'N'}.")
print(f"- Token={'Y' if token else 'N'}.")
this_init_string = bundle_path_to_init_string(bundle_path, keyspace, token)
print("=> Export command:")
print(f'export ASTRA_DB_INIT_STRING="{this_init_string}"')
if __name__ == "__main__":
create_init_string_utility()
|