# Copyright 2021 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import re
import test_util
import time
from absl import app
from selenium import webdriver
from pywinauto.application import Application

UnsafePageLink = "http://testsafebrowsing.appspot.com/s/malware.html"
UnsafeDownloadLink = "http://testsafebrowsing.appspot.com/s/badrep.exe"


def visit(window, url):
  """Visit a specific URL through pywinauto.Application.

  SafeBrowsing intercepts HTTP requests & hangs WebDriver.get(), which prevents
  us from getting the page source. Using pywinauto to visit the pages instead.
  """
  window.Edit.set_edit_text(url).type_keys("%{ENTER}")
  time.sleep(10)


def main(argv):
  exclude_switches = ["disable-background-networking"]
  chrome_options = webdriver.ChromeOptions()
  chrome_options.add_experimental_option("excludeSwitches", exclude_switches)

  driver = test_util.create_chrome_webdriver(chrome_options=chrome_options)

  try:
    app = Application(backend="uia")
    app.connect(title_re='.*Chrome|.*Chromium')
    window = app.top_window()

    # Wait for Chrome to download SafeBrowsing lists in the background.
    # There's no trigger to force this operation or synchronize on it, but quick
    # experiments have shown 3-4 minutes in most cases, so 5 should be plenty.
    time.sleep(60 * 5)

    # Verify Policy status legend in chrome://policy page
    policy_url = "chrome://policy"
    driver.get(policy_url)
    driver.find_element_by_id('reload-policies').click
    deviceId = driver.find_element_by_class_name(
        'machine-enrollment-device-id').text

    visit(window, UnsafePageLink)

    visit(window, UnsafeDownloadLink)
    print('\nDeviceId:' + deviceId.strip())

  except Exception as error:
    print(error)
  finally:
    driver.quit()


if __name__ == '__main__':
  app.run(main)
