Example #1
0
 def test_write_to_neptune_lpg_graph(self):
     endpoint = NeptuneEndpoint(host="host", port=5555, region="us-east-1", ssl=False)
     client = AltimeterNeptuneClient(0, endpoint)
     # An exception is expected here since there is no graph to write to
     with self.assertRaises(NeptuneLoadGraphException):
         client.write_to_neptune_lpg({"vertices": [
             {"~id": "123", "~label": "test"}],
             "edges" : []}, "")
 def test_write_to_neptune_lpg_no_graph(self):
     endpoint = NeptuneEndpoint(host="host",
                                port=5555,
                                region="us-east-1",
                                ssl=False)
     client = AltimeterNeptuneClient(0, endpoint)
     with self.assertRaises(NeptuneNoGraphsFoundException):
         client.write_to_neptune_lpg({}, "")
def aws2neptune_lpg(scan_id: str, config: AWSConfig,
                    muxer: AWSScanMuxer) -> None:
    """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)

    aws_resource_region_mapping_repo = build_aws_resource_region_mapping_repo(
        global_region_whitelist=config.scan.regions,
        preferred_account_scan_regions=config.scan.
        preferred_account_scan_regions,
        services_regions_json_url=config.services_regions_json_url,
    )

    logger.info(
        AWSLogEvents.ScanConfigured,
        config=str(config),
        reader=str(artifact_reader.__class__),
        writer=str(artifact_writer.__class__),
    )
    print("Beginning AWS Account Scan")

    _, graph_set = run_scan(
        muxer=muxer,
        config=config,
        aws_resource_region_mapping_repo=aws_resource_region_mapping_repo,
        artifact_writer=artifact_writer,
        artifact_reader=artifact_reader,
    )
    print("AWS Account Scan Complete. Beginning write to Amazon Neptune.")
    logger.info(LogEvent.NeptuneGremlinWriteStart)
    graph = graph_set.to_neptune_lpg(scan_id)
    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,
        ssl=bool(config.neptune.ssl),
        auth_mode=str(config.neptune.auth_mode),
    )
    neptune_client = AltimeterNeptuneClient(max_age_min=1440,
                                            neptune_endpoint=endpoint)
    neptune_client.write_to_neptune_lpg(graph, scan_id)
    logger.info(LogEvent.NeptuneGremlinWriteEnd)
    print("Write to Amazon Neptune Complete")