// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/safe_browsing/content/renderer/phishing_classifier/content_phishing_image_embedder_delegate.h"

#include <memory>
#include <optional>
#include <utility>

#include "base/metrics/histogram_macros.h"
#include "components/safe_browsing/content/common/safe_browsing.mojom-shared.h"
#include "components/safe_browsing/content/renderer/phishing_classifier/content_phishing_image_embedder.h"
#include "components/safe_browsing/core/common/phishing_classifier/scorer.h"
#include "components/safe_browsing/core/common/proto/csd.pb.h"
#include "content/public/renderer/render_frame.h"
#include "mojo/public/cpp/base/proto_wrapper.h"
#include "third_party/blink/public/common/associated_interfaces/associated_interface_registry.h"
#include "third_party/blink/public/web/web_document.h"
#include "third_party/blink/public/web/web_local_frame.h"

namespace safe_browsing {

namespace {

GURL StripRef(const GURL& url) {
  GURL::Replacements replacements;
  replacements.ClearRef();
  return url.ReplaceComponents(replacements);
}

}  // namespace

ContentPhishingImageEmbedderDelegate::ContentPhishingImageEmbedderDelegate(
    content::RenderFrame* render_frame)
    : content::RenderFrameObserver(render_frame),
      image_embedder_(
          std::make_unique<ContentPhishingImageEmbedder>(render_frame)),
      is_image_embedding_(false),
      is_image_embedding_running_(false) {
  render_frame->GetAssociatedInterfaceRegistry()
      ->AddInterface<mojom::PhishingImageEmbedderDetector>(base::BindRepeating(
          &ContentPhishingImageEmbedderDelegate::PhishingImageEmbedderReceiver,
          base::Unretained(this)));

  model_change_observation_.Observe(ScorerStorage::GetInstance());
}

ContentPhishingImageEmbedderDelegate::~ContentPhishingImageEmbedderDelegate() {
  CancelPendingImageEmbedding(kShutdown);
}

// static
ContentPhishingImageEmbedderDelegate*
ContentPhishingImageEmbedderDelegate::Create(
    content::RenderFrame* render_frame) {
  return new ContentPhishingImageEmbedderDelegate(render_frame);
}

void ContentPhishingImageEmbedderDelegate::PhishingImageEmbedderReceiver(
    mojo::PendingAssociatedReceiver<mojom::PhishingImageEmbedderDetector>
        receiver) {
  phishing_image_embedder_receiver_.reset();
  phishing_image_embedder_receiver_.Bind(std::move(receiver));
}

void ContentPhishingImageEmbedderDelegate::StartImageEmbedding(
    const GURL& url,
    bool can_extract_visual_features,
    StartImageEmbeddingCallback callback) {
  RecordEvent(SBPhishingImageEmbedderEvent::kPhishingImageEmbeddingRequested);
  if (image_embedding_callback_) {
    std::move(image_embedding_callback_)
        .Run(mojom::PhishingImageEmbeddingResult::kCancelled, std::nullopt,
             std::nullopt);
    CancelPendingImageEmbedding(
        CancelImageEmbeddingReason::kNewRequestFromBrowser);
  }
  is_image_embedding_running_ = true;
  last_url_received_from_browser_ = StripRef(url);
  image_embedding_callback_ = std::move(callback);

  MaybeStartImageEmbedding(can_extract_visual_features);
}

void ContentPhishingImageEmbedderDelegate::DidCommitProvisionalLoad(
    ui::PageTransition transition) {
  blink::WebLocalFrame* frame = render_frame()->GetWebFrame();
  // A new page is starting to load, so cancel image embedding.
  CancelPendingImageEmbedding(kNavigateAway);
  if (!frame->Parent()) {
    last_main_frame_transition_ = transition;
  }
}

void ContentPhishingImageEmbedderDelegate::DidFinishSameDocumentNavigation() {
  CancelPendingImageEmbedding(kNavigateWithinPage);
}

bool ContentPhishingImageEmbedderDelegate::is_ready() const {
  return image_embedder_->is_ready();
}

void ContentPhishingImageEmbedderDelegate::MaybeStartImageEmbedding(
    bool can_extract_visual_features) {
  // We can begin the image embedding process when the following conditions are
  // met:
  //  1. A Scorer has been created
  //  2. The browser has sent a StartImageEmbedding message for the
  //     current toplevel URL.
  //  3. The load is a new navigation (not a session history navigation).
  //  4. The toplevel URL has not already been processed.
  if (!image_embedder_->is_ready()) {
    is_image_embedding_running_ = false;
    if (!image_embedding_callback_.is_null()) {
      std::move(image_embedding_callback_)
          .Run(mojom::PhishingImageEmbeddingResult::kImageEmbedderNotReady,
               std::nullopt, std::nullopt);
    }
    return;
  }

  if (last_main_frame_transition_ & ui::PAGE_TRANSITION_FORWARD_BACK) {
    // Skip loads from session history navigation.  However, update the
    // last URL sent to the image embedder, so that we'll properly detect
    // same-document navigations.
    last_url_sent_to_image_embedder_ = last_finished_load_url_;
    is_image_embedding_running_ = false;
    if (!image_embedding_callback_.is_null()) {
      std::move(image_embedding_callback_)
          .Run(mojom::PhishingImageEmbeddingResult::kForwardBackTransition,
               std::nullopt, std::nullopt);
    }
    return;
  }

  GURL stripped_last_load_url(StripRef(last_finished_load_url_));

  if (last_url_received_from_browser_ != stripped_last_load_url) {
    RecordEvent(
        SBPhishingImageEmbedderEvent::kUrlShouldNotBeUsedForImageEmbedding);
    // The browser has not yet confirmed that this URL should be used for image
    // embedding, so defer image embedding for now.  Note: the ref does not
    // affect any of the browser's pre image embedding checks, so we don't
    // require it to match.
    return;
  }

  last_url_sent_to_image_embedder_ = last_finished_load_url_;
  is_image_embedding_ = true;
  image_embedder_->BeginImageEmbedding(
      can_extract_visual_features,
      base::BindOnce(&ContentPhishingImageEmbedderDelegate::ImageEmbeddingDone,
                     base::Unretained(this)));
}

void ContentPhishingImageEmbedderDelegate::PageCaptured(
    bool preliminary_capture) {
  RecordEvent(SBPhishingImageEmbedderEvent::kPageTextCaptured);

  if (preliminary_capture) {
    return;
  }
  // Make sure there's no image embedding in progress.  We don't want to swap
  // out the page text string from underneath the term feature extractor.
  //
  // Note: Currently, if the url hasn't changed, we won't restart the image
  // embedding process in this case.  We may want to adjust this.
  CancelPendingImageEmbedding(kPageRecaptured);
  last_finished_load_url_ = render_frame()->GetWebFrame()->GetDocument().Url();

  GURL stripped_last_load_url(StripRef(last_finished_load_url_));
  // Check if toplevel URL has changed.
  if (stripped_last_load_url == StripRef(last_url_sent_to_image_embedder_)) {
    return;
  }
  RecordEvent(SBPhishingImageEmbedderEvent::kImageEmbeddingBegins);
  MaybeStartImageEmbedding();
}

void ContentPhishingImageEmbedderDelegate::CancelPendingImageEmbedding(
    CancelImageEmbeddingReason reason) {
  if (is_image_embedding_) {
    UMA_HISTOGRAM_ENUMERATION("SBClientPhishing.CancelImageEmbeddingReason",
                              reason, CancelImageEmbeddingReason::kMaxValue);
    is_image_embedding_ = false;
    image_embedder_->CancelPendingImageEmbedding();
  }
}

void ContentPhishingImageEmbedderDelegate::ImageEmbeddingDone(
    ContentPhishingImageEmbedder::Result result,
    const ImageFeatureEmbedding& image_feature_embedding,
    const VisualFeatures& visual_features) {
  RecordEvent(SBPhishingImageEmbedderEvent::kImageEmbeddingComplete);
  is_image_embedding_running_ = false;
  if (image_embedding_callback_.is_null()) {
    RecordEvent(
        SBPhishingImageEmbedderEvent::kImageEmbeddingCallbackEmptyOnCompletion);
    return;
  }

  RecordEvent(SBPhishingImageEmbedderEvent::kImageEmbeddingRequestResponded);
  mojom::PhishingImageEmbeddingResult final_result =
      mojom::PhishingImageEmbeddingResult::kSuccess;
  switch (result) {
    case ContentPhishingImageEmbedder::Result::kSuccess:
      if (image_feature_embedding.embedding_value_size() == 0) {
        final_result = mojom::PhishingImageEmbeddingResult::kFailed;
      }
      break;
    case ContentPhishingImageEmbedder::Result::kInvalidURLFormatRequest:
      final_result =
          mojom::PhishingImageEmbeddingResult::kInvalidURLFormatRequest;
      break;
    case ContentPhishingImageEmbedder::Result::kInvalidDocumentLoader:
      final_result =
          mojom::PhishingImageEmbeddingResult::kInvalidDocumentLoader;
      break;
    case ContentPhishingImageEmbedder::Result::kVisualExtractionFailed:
      final_result = mojom::PhishingImageEmbeddingResult::kFailed;
      break;
  }

  if (final_result == mojom::PhishingImageEmbeddingResult::kSuccess) {
    std::move(image_embedding_callback_)
        .Run(final_result, mojo_base::ProtoWrapper(image_feature_embedding),
             mojo_base::ProtoWrapper(visual_features));
  } else {
    std::move(image_embedding_callback_)
        .Run(final_result, std::nullopt, std::nullopt);
  }
}

void ContentPhishingImageEmbedderDelegate::RecordEvent(
    SBPhishingImageEmbedderEvent event) {
  UMA_HISTOGRAM_ENUMERATION("SBClientPhishing.ImageEmbedder.Event2", event);
}

void ContentPhishingImageEmbedderDelegate::OnDestruct() {
  if (is_image_embedding_running_) {
    RecordEvent(
        SBPhishingImageEmbedderEvent::kDestructedBeforeImageEmbeddingDone);
  }
  delete this;
}

void ContentPhishingImageEmbedderDelegate::OnScorerChanged() {
  Scorer* scorer = ScorerStorage::GetInstance()->GetScorer();

  if (!scorer) {
    // If the scorer is reset, we should clear pending image embedding if there
    // is one going on, which is checked by the function below.
    CancelPendingImageEmbedding(kScorerCleared);
    return;
  }

  if (is_image_embedding_) {
    CancelPendingImageEmbedding(kNewPhishingScorer);
  }
}

}  // namespace safe_browsing