/*
 * Copyright 2020 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
 *
 * http://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.
 */

package com.example.containeranalysis;

import com.google.cloud.devtools.containeranalysis.v1.ContainerAnalysisClient;
import com.google.cloud.functions.BackgroundFunction;
import com.google.cloud.functions.Context;
import com.google.gson.Gson;
import io.grafeas.v1.GrafeasClient;
import io.grafeas.v1.NoteKind;
import io.grafeas.v1.Occurrence;
import io.grafeas.v1.Severity;
import io.grafeas.v1.VulnerabilityOccurrence;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

public class VulnerabilityFunction implements BackgroundFunction<PubSubMessage> {

  private static final Gson gson = new Gson();

  private static final Logger logger = Logger.getLogger(VulnerabilityFunction.class.getName());

  private final ContainerAnalysisClient client;
  private final GrafeasClient grafeasClient;

  public VulnerabilityFunction() throws IOException {
    this(ContainerAnalysisClient.create());
  }

  VulnerabilityFunction(ContainerAnalysisClient client) throws IOException {
    this.client = client;
    this.grafeasClient = client.getGrafeasClient();

    // When the funciton is being terminated, shutdown the clients gracefully.
    Runtime.getRuntime().addShutdownHook(new Thread(() -> {
      grafeasClient.shutdownNow();
      client.shutdownNow();
      try {
        grafeasClient.awaitTermination(5, TimeUnit.SECONDS);
      } catch (InterruptedException e) {
        // safely ignore
      }
      try {
        client.awaitTermination(5, TimeUnit.SECONDS);
      } catch (InterruptedException e) {
        // safely ignore
      }
    }));
  }

  @Override
  public void accept(PubSubMessage payload, Context context) {
    String json = new String(Base64.getDecoder().decode(payload.getData()), StandardCharsets.UTF_8);
    OccurrenceNotification notification = gson.fromJson(json, OccurrenceNotification.class);

    // Retrieve the occurrence detials from the notification
    // https://cloud.google.com/container-registry/docs/reference/rest/v1/projects.occurrences#Occurrence
    Occurrence occurrence = grafeasClient.getOccurrence(notification.getName());

    // If the occurence is a vulnerability, output the log line base on severity
    if (NoteKind.VULNERABILITY.equals(occurrence.getKind())) {
      VulnerabilityOccurrence vulnerability = occurrence.getVulnerability();
      if (vulnerability.getSeverity().getNumber() >= Severity.HIGH_VALUE) {
        logger.warning(
            String.format(
                "Image: %s, CVE: %s, Severity: %s",
                occurrence.getResourceUri(),
                vulnerability.getShortDescription(),
                vulnerability.getSeverity()));
      }
    }
  }
}
