# Copyright 2023 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import argparse

import json

import logging

from apache_beam import (

CombineFn,

CombineGlobally,

DoFn,

io,

ParDo,

Pipeline,

WindowInto,

)

from apache_beam.error import PipelineError

from apache_beam.options.pipeline_options import GoogleCloudOptions, PipelineOptions

from apache_beam.transforms.window import FixedWindows

from google.cloud import dlp_v2, logging_v2

# For more details about info types format, please see

# https://cloud.google.com/dlp/docs/reference/rest/v2/InspectConfig

INSPECT_CFG = {"info_types": [{"name": "US_SOCIAL_SECURITY_NUMBER"}]}

# For more details about transformation format, please see

# https://cloud.google.com/dlp/docs/reference/rest/v2/projects.deidentifyTemplates#DeidentifyTemplate.InfoTypeTransformations

REDACTION_CFG = {

"info_type_transformations": {

"transformations": [

{

"primitive_transformation": {

"character_mask_config": {"masking_character": "#"}

}

}

]

}

}

class PayloadAsJson(DoFn):

"""Convert PubSub message payload to UTF-8 and return as JSON"""

def process(self, element):

yield json.loads(element.decode("utf-8"))

class BatchPayloads(CombineFn):

"""Collect all items in the windowed collection into single batch"""

def create_accumulator(self):

return []

def add_input(self, accumulator, input):

accumulator.append(input)

return accumulator

def merge_accumulators(self, accumulators):

merged = [item for accumulator in accumulators for item in accumulator]

return merged

def extract_output(self, accumulator):

return accumulator

class LogRedaction(DoFn):

"""Apply inspection and redaction to textPayload field of log entries"""

def __init__(self, region, project_id: str):

self.project_id = project_id

self.region = region

self.dlp_client = None

def _log_to_row(self, entry):

# Make `Row` from `textPayload`. For more details on the row, please see

# https://cloud.google.com/dlp/docs/reference/rest/v2/ContentItem#Row

payload = entry.get("textPayload", "")

return {"values": [{"string_value": payload}]}

def setup(self):

"""Initialize DLP client"""

if self.dlp_client:

return

self.dlp_client = dlp_v2.DlpServiceClient()

if not self.dlp_client:

logging.error("Cannot create Google DLP Client")

raise PipelineError("Cannot create Google DLP Client")

def process(self, logs):

# Construct the `table`. For more details on the table schema, please see

# https://cloud.google.com/dlp/docs/reference/rest/v2/ContentItem#Table

table = {

"table": {

"headers": [{"name": "textPayload"}],

"rows": map(self._log_to_row, logs),

}

}

response = self.dlp_client.deidentify_content(

request={

"parent": f"projects/{self.project_id}/locations/{self.region}",

"inspect_config": INSPECT_CFG,

"deidentify_config": REDACTION_CFG,

"item": table,

}

)

# replace payload with redacted version

modified_logs = []

for index, log in enumerate(logs):

log["textPayload"] = response.item.table.rows[index].values[0].string_value

# you may consider changing insert ID if the project already has a copy

# of this log (e.g. log['insertId'] = 'deid-' + log['insertId'])

# For more details about insert ID, please see:

# https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry#FIELDS.insert_id

modified_logs.append(log)

yield modified_logs

class IngestLogs(DoFn):

"""Ingest payloads into destination log"""

def __init__(self, destination_log_name):

self.destination_log_name = destination_log_name

self.logger = None

def _replace_log_name(self, entry):

# updates log name in the entry to logger name

entry["logName"] = self.logger.name

return entry

def setup(self):

# initialize logging client

if self.logger:

return

logging_client = logging_v2.Client()

if not logging_client:

logging.error("Cannot create Google Logging Client")

raise PipelineError("Cannot create Google Logging Client")

self.logger = logging_client.logger(self.destination_log_name)

if not self.logger:

logging.error("Google client library cannot create Logger object")

raise PipelineError("Google client library cannot create Logger object")

def process(self, element):

if self.logger:

logs = list(map(self._replace_log_name, element))

self.logger.client.logging_api.write_entries(logs)

yield logs

def run(

pubsub_subscription: str,

destination_log_name: str,

window_size: float,

pipeline_args: list[str] = None,

) -> None:

"""Runs Dataflow pipeline"""

pipeline_options = PipelineOptions(

pipeline_args, streaming=True, save_main_session=True

)

region = "us-central1"

try:

region = pipeline_options.view_as(GoogleCloudOptions).region

except AttributeError:

pass

pipeline = Pipeline(options=pipeline_options)

_ = (

pipeline

| "Read log entries from Pub/Sub"

>> io.ReadFromPubSub(subscription=pubsub_subscription)

| "Convert log entry payload to Json" >> ParDo(PayloadAsJson())

| "Aggregate payloads in fixed time intervals"

>> WindowInto(FixedWindows(window_size))

# Optimize Google API consumption and avoid possible throttling

# by calling APIs for batched data and not per each element

| "Batch aggregated payloads"

>> CombineGlobally(BatchPayloads()).without_defaults()

| "Redact SSN info from logs"

>> ParDo(LogRedaction(region, destination_log_name.split("/")[1]))

| "Ingest to output log" >> ParDo(IngestLogs(destination_log_name))

)

pipeline.run()

if __name__ == "__main__":

logging.getLogger().setLevel(logging.INFO)

parser = argparse.ArgumentParser()

parser.add_argument(

"--pubsub_subscription",

help="The Cloud Pub/Sub subscription to read from in the format "

'"projects/<PROJECT_ID>/subscription/<SUBSCRIPTION_ID>".',

)

parser.add_argument(

"--destination_log_name",

help="The log name to ingest log entries in the format "

'"projects/<PROJECT_ID>/logs/<LOG_ID>".',

)

parser.add_argument(

"--window_size",

type=float,

default=60.0,

help="Output file's window size in seconds.",

)

known_args, pipeline_args = parser.parse_known_args()

run(

known_args.pubsub_subscription,

known_args.destination_log_name,

known_args.window_size,

pipeline_args,

)