Example #1
0
def lambda_handler(event, context):
    rdf_bucket = event["Records"][0]["s3"]["bucket"]["name"]
    rdf_key = urllib.parse.unquote(event["Records"][0]["s3"]["object"]["key"])

    neptune_host = get_required_lambda_env_var("NEPTUNE_HOST")
    neptune_port = get_required_lambda_env_var("NEPTUNE_PORT")
    neptune_region = get_required_lambda_env_var("NEPTUNE_REGION")
    neptune_load_iam_role_arn = get_required_lambda_env_var(
        "NEPTUNE_LOAD_IAM_ROLE_ARN")
    on_success_sns_topic_arn = get_required_lambda_env_var(
        "ON_SUCCESS_SNS_TOPIC_ARN")

    endpoint = NeptuneEndpoint(host=neptune_host,
                               port=neptune_port,
                               region=neptune_region)
    neptune_client = AltimeterNeptuneClient(max_age_min=1440,
                                            neptune_endpoint=endpoint)
    graph_metadata = neptune_client.load_graph(
        bucket=rdf_bucket,
        key=rdf_key,
        load_iam_role_arn=neptune_load_iam_role_arn)

    logger = Logger()
    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=on_success_sns_topic_arn,
                       MessageStructure="json",
                       Message=json.dumps(message_dict))
    logger.info(event=LogEvent.GraphLoadedSNSNotificationEnd)
Example #2
0
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)