def from_artifact_path(cls: Type["ArtifactReader"],
artifact_path: str) -> "ArtifactReader":
"""Create an ArtifactReader based on an artifact path. This either returns a
FileArtifactReader or an S3ArtifactReader depending on the value of artifact_path"""
if is_s3_uri(artifact_path):
parse_s3_uri(artifact_path)
return S3ArtifactReader()
return FileArtifactReader()
def read_json(self, path: str) -> Dict[str, Any]:
"""Read a json artifact
Args:
path: s3 uri to artifact. s3://bucket/key/path
Returns:
artifact content
"""
bucket, key = parse_s3_uri(path)
if key is None:
raise ValueError(f"Unable to read from s3 uri missing key: {path}")
session = boto3.Session()
s3_client = session.client("s3")
logger = Logger()
with logger.bind(bucket=bucket, key=key):
with io.BytesIO() as artifact_bytes_buf:
logger.info(event=LogEvent.ReadFromS3Start)
s3_client.download_fileobj(bucket, key, artifact_bytes_buf)
artifact_bytes_buf.flush()
artifact_bytes_buf.seek(0)
artifact_bytes = artifact_bytes_buf.read()
logger.info(event=LogEvent.ReadFromS3End)
artifact_str = artifact_bytes.decode("utf-8")
artifact_dict = json.loads(artifact_str)
return artifact_dict
def __post_init__(self) -> None:
if (self.scan.single_account_mode and not self.scan.scan_sub_accounts
and self.access.accessor.multi_hop_accessors):
raise InvalidConfigException(
"Accessor config not supported for single account mode")
if is_s3_uri(self.artifact_path):
bucket, key_prefix = parse_s3_uri(self.artifact_path)
if key_prefix is not None:
raise InvalidConfigException(
f"S3 artifact_path should be s3://<bucket>, no key - got {self.artifact_path}"
)
def from_artifact_path(cls: Type["ArtifactReader"],
artifact_path: str) -> "ArtifactReader":
"""Create an ArtifactReader based on an artifact path. This either returns a
FileArtifactReader or an S3ArtifactReader depending on the value of artifact_path"""
if is_s3_uri(artifact_path):
_, key_prefix = parse_s3_uri(artifact_path)
if key_prefix is not None:
raise ValueError(
f"S3 artifact path should be s3://<bucket>, no key - got {artifact_path}"
)
return S3ArtifactReader()
return FileArtifactReader()
def from_artifact_path(cls: Type["ArtifactWriter"], artifact_path: str,
scan_id: str) -> "ArtifactWriter":
"""Create an ArtifactWriter based on an artifact path. This either returns a FileArtifactWriter
or an S3ArtifactWriter depending on the value of artifact_path"""
if is_s3_uri(artifact_path):
bucket, key_prefix = parse_s3_uri(artifact_path)
if key_prefix is not None:
raise ValueError(
f"S3 artifact path should be s3://<bucket>, no key - got {artifact_path}"
)
return S3ArtifactWriter(bucket=bucket, key_prefix=scan_id)
return FileArtifactWriter(scan_id=scan_id,
output_dir=Path(artifact_path))
def from_artifact_path(cls: Type["ArtifactWriter"], artifact_path: str,
scan_id: str) -> "ArtifactWriter":
"""Create an ArtifactWriter based on an artifact path. This either returns a FileArtifactWriter
or an S3ArtifactWriter depending on the value of artifact_path"""
if is_s3_uri(artifact_path):
bucket, s3_uri_key_prefix = parse_s3_uri(artifact_path)
if s3_uri_key_prefix is not None:
key_prefix = "/".join((s3_uri_key_prefix, scan_id))
else:
key_prefix = scan_id
return S3ArtifactWriter(bucket=bucket, key_prefix=key_prefix)
return FileArtifactWriter(scan_id=scan_id,
output_dir=Path(artifact_path))
def from_s3(cls: Type["Config"], s3_uri: str) -> "Config":
"""Load a Config from an s3 object"""
bucket, key = parse_s3_uri(s3_uri)
s3_client = boto3.client("s3")
resp = s3_client.get_object(
Bucket=bucket,
Key=key,
)
config_str = resp["Body"].read().decode("utf-8")
config_dict = dict(toml.loads(config_str))
try:
return cls.from_dict(config_dict)
except InvalidConfigException as ice:
raise InvalidConfigException(
f"Error in conf file {s3_uri}: {str(ice)}")
def aws2n(scan_id: str, config: Config, muxer: AWSScanMuxer,
load_neptune: bool) -> AWS2NResult:
"""Scan AWS resources to json, convert to RDF and load into Neptune
if config.neptune is defined"""
artifact_reader = ArtifactReader.from_artifact_path(config.artifact_path)
artifact_writer = ArtifactWriter.from_artifact_path(
artifact_path=config.artifact_path, scan_id=scan_id)
logger = Logger()
logger.info(
AWSLogEvents.ScanConfigured,
config=str(config),
reader=str(artifact_reader.__class__),
writer=str(artifact_writer.__class__),
)
scan_manifest, graph_set = run_scan(
muxer=muxer,
config=config,
artifact_writer=artifact_writer,
artifact_reader=artifact_reader,
)
json_path = scan_manifest.master_artifact
rdf_path = artifact_writer.write_graph_set(name="master",
graph_set=graph_set,
compression=GZIP)
graph_metadata = None
if load_neptune:
if config.neptune is None:
raise Exception(
"Can not load to Neptune because config.neptune is empty.")
endpoint = NeptuneEndpoint(host=config.neptune.host,
port=config.neptune.port,
region=config.neptune.region)
neptune_client = AltimeterNeptuneClient(max_age_min=1440,
neptune_endpoint=endpoint)
rdf_bucket, rdf_key = parse_s3_uri(rdf_path)
if rdf_key is None:
raise Exception(f"Invalid rdf s3 path {rdf_path}")
graph_metadata = neptune_client.load_graph(
bucket=rdf_bucket,
key=rdf_key,
load_iam_role_arn=str(config.neptune.iam_role_arn))
logger.info(event=LogEvent.GraphLoadedSNSNotificationStart)
sns_client = boto3.client("sns")
message_dict = {
"uri": graph_metadata.uri,
"name": graph_metadata.name,
"version": graph_metadata.version,
"start_time": graph_metadata.start_time,
"end_time": graph_metadata.end_time,
"neptune_endpoint": endpoint.get_endpoint_str(),
}
message_dict["default"] = json.dumps(message_dict)
sns_client.publish(
TopicArn=config.neptune.graph_load_sns_topic_arn,
MessageStructure="json",
Message=json.dumps(message_dict),
)
logger.info(event=LogEvent.GraphLoadedSNSNotificationEnd)
return AWS2NResult(json_path=json_path,
rdf_path=rdf_path,
graph_metadata=graph_metadata)
def __init__(self, **data: Any):
super().__init__(**data)
if is_s3_uri(self.artifact_path):
parse_s3_uri(self.artifact_path)
def test_bad_key(self):
uri = "s3://bucket/key//goo"
with self.assertRaises(InvalidS3URIException):
parse_s3_uri(uri)
def test_without_key(self):
uri = "s3://bucket/"
bucket, key = parse_s3_uri(uri)
self.assertIsNone(key)
self.assertEqual(bucket, "bucket")
def test_without_bucket(self):
uri = "s3:///key/path"
with self.assertRaises(InvalidS3URIException):
parse_s3_uri(uri)
def test_valid_long_key(self):
uri = "s3://bucket/key/foo/boo/goo"
bucket, key = parse_s3_uri(uri)
self.assertEqual(bucket, "bucket")
self.assertEqual(key, "key/foo/boo/goo")
def test_valid_simple_key(self):
uri = "s3://bucket/key"
bucket, key = parse_s3_uri(uri)
self.assertEqual(bucket, "bucket")
self.assertEqual(key, "key")
def test_invalid_uri(self):
uri = "bucket/key"
with self.assertRaises(InvalidS3URIException):
parse_s3_uri(uri)