Automatically Update ChromeDriver to Match Google Chrome Version for Selenium Scripts with HTTP/2 Streaming in Python

"""
    Description:

    Automatically updates the chromedriver.

    Requirements:

    Need to have sudo command permissions without password. To do this, run:
    echo "$USER ALL=(ALL:ALL) NOPASSWD: ALL" | sudo tee /etc/sudoers.d/dont-prompt-$USER-for-sudo-password
    You will be prompted for the password one last time.

    python3 -m pip install --upgrade pip
    pip install tqdm httpx[http2]


    Author:           Marcelo Facio Palin
    Created:          2023-10-31
"""
import subprocess
import platform
import httpx
from tqdm import tqdm
import os
from pathlib import Path
import zipfile
import shutil

def get_chrome_version():
    try:
        version = subprocess.check_output(["google-chrome", "--version"]).decode("utf-8").split(" ")[2].strip()
        return version
    except Exception as e:
        print("Could not determine Google Chrome version. Make sure Google Chrome is installed and in the PATH.")
        print(e)
        return None

def get_platform():
    system = platform.system().lower()
    if system == "linux":
        return "linux64"
    else:
        return system

def get_chromedriver_version(chromedriver_path):
    try:
        version = subprocess.check_output([chromedriver_path, "--version"]).decode("utf-8").strip()
        return version
    except Exception as e:
        print(f"Could not determine ChromeDriver version at {chromedriver_path}. Error: {e}")
        return None

def download_chromedriver(chrome_version, platform):
    url = "https://googlechromelabs.github.io/chrome-for-testing/last-known-good-versions-with-downloads.json"
    response = httpx.get(url)
    if response.status_code != 200:
        print("Could not get download information. Status code:", response.status_code)
        return

    data = response.json()
    for channel, channel_data in data["channels"].items():
        if channel_data["version"] == chrome_version:
            for download in channel_data["downloads"]["chromedriver"]:
                if download["platform"] == platform:
                    download_url = download["url"]
                    print(f"Download link for ChromeDriver version {chrome_version} for {platform}: {download_url}")
                    return download_url

    print(f"Could not find ChromeDriver for version {chrome_version} and platform {platform}.")
    return None

def download_file(url, filename):
    with httpx.Client(http2=True) as client:
        with client.stream("GET", url) as response:
            if response.status_code != 200:
                print("Could not download the file. Status code:", response.status_code)
                return

            file_size = int(response.headers.get("Content-Length", 0))
            progress_bar = tqdm(total=file_size, unit="B", unit_scale=True, desc=filename)

            with open(filename, "wb") as file:
                for chunk in response.iter_bytes(1024):
                    progress_bar.update(len(chunk))
                    file.write(chunk)

            progress_bar.close()

    print(f"File saved at: {os.path.abspath(filename)}")

def install_chromedriver(zip_filename):
    extract_dir = "chromedriver_extract"
    with zipfile.ZipFile(zip_filename, 'r') as zip_ref:
        zip_ref.extractall(extract_dir)

    chromedriver_path = None
    for root, dirs, files in os.walk(extract_dir):
        if "chromedriver" in files:
            chromedriver_path = os.path.join(root, "chromedriver")
            break

    if chromedriver_path is None:
        print("Could not find the executable file of chromedriver.")
        return

    # Makes the file executable
    os.chmod(chromedriver_path, 0o755)

    # Installs chromedriver in /usr/local/bin
    try:
        subprocess.run(["sudo", "cp", chromedriver_path, "/usr/local/bin/chromedriver"], check=True)
    except subprocess.CalledProcessError as e:
        print("Could not install chromedriver in /usr/local/bin. Error: ", e)

    dir_path = Path("/opt/chromedriver")
    if not dir_path.exists():
        try:
            subprocess.run(["sudo", "mkdir", str(dir_path)], check=True)
        except subprocess.CalledProcessError as e:
            print("Could not create the directory /opt/chromedriver. Error: ", e)

    # Installs chromedriver in /opt/chromedriver
    try:
        subprocess.run(["sudo", "cp", chromedriver_path, "/opt/chromedriver/chromedriver"], check=True)
    except subprocess.CalledProcessError as e:
        print("Could not install chromedriver in /opt/chromedriver. Error: ", e)

    # Removes the extraction folder
    shutil.rmtree(extract_dir)

    print("Chromedriver successfully installed!")


def main():
    chrome_version = get_chrome_version()
    platform = get_platform()
    if chrome_version and platform:
        download_url = download_chromedriver(chrome_version, platform)

        if download_url:
            zip_filename = f"chromedriver_{chrome_version}_{platform}.zip"
            download_file(download_url, zip_filename)

            install_chromedriver(zip_filename)

            print("Download and installation successful!")
            print("Google Chrome version: ", chrome_version)

            usr_bin_version = get_chromedriver_version("/usr/local/bin/chromedriver")
            if usr_bin_version:
                print("ChromeDriver version in /usr/local/bin: ", usr_bin_version)

            opt_chromedriver_version = get_chromedriver_version("/opt/chromedriver/chromedriver")
            if opt_chromedriver_version:
                print("ChromeDriver version in /opt/chromedriver/chromedriver: ", opt_chromedriver_version)

            print("Platform: ", platform)
        else:
            print("The download was not successful.")

if __name__ == "__main__":
    main()