A Python Script to Automatically Update ChromeDriver

I wrote a script to install the latest version of ChromeDriver and the compatible test version of Chrome. (Supports version 115 and above)

The script will prompt for the sudo password.

#!/usr/bin/env python3
"""
Script to automatically update ChromeDriver and Chrome test version
"""
import getpass
import shutil
import subprocess
import pathlib
from contextlib import contextmanager

import requests
import os

USER_AGENT = '@ytyng/upgrade_chromedriver'

default_request_headers = {
    'User-Agent': USER_AGENT
}

DOWNLOAD_DIR = os.path.expanduser('~/Downloads')

@contextmanager
def set_directory(path: pathlib.Path):
    """Sets the cwd within the context

    Args:
        path (Path): The path to the cwd

    Yields:
        None
    """

    origin = pathlib.Path().absolute()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(origin)

def get_cpu_architecture() -> str:
    """
    Get the CPU architecture
    """
    return subprocess.run(
        ['uname', '-m'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    ).stdout.decode('utf-8').strip()

def get_current_machine_platform():
    """
    Returns a string that matches the platform in JSON.
    Besides mac-arm64 and mac-x64, there are also linux64, win32, and win64,
    but currently only mac is supported.
    """
    a = get_cpu_architecture()
    if a == 'arm64':
        return 'mac-arm64'
    else:
        return 'mac-x64'

json_url = 'https://googlechromelabs.github.io/chrome-for-testing/known-good-versions-with-downloads.json'

_rdata = requests.get(json_url, headers=default_request_headers).json()

# Filter to only include versions that have both chrome and chromedriver
_versions = [v for v in _rdata['versions'] if
             v['downloads'].get('chromedriver') and v['downloads'].get(
                 'chrome')]

_versions = sorted(_versions,
                   key=lambda x: tuple(map(int, x['version'].split('.'))),
                   reverse=True)

def download_latest_version(version_dict, app_name: str):
    """
    Download the file
    :param app_name: str 'chrome' or 'chromedriver'
    """
    print(f'Downloading {app_name}...')
    _platform = get_current_machine_platform()
    target_url = [d for d in version_dict['downloads'][app_name] if
                  d['platform'] == _platform][0]['url']
    file_binary = requests.get(target_url, headers=default_request_headers)

    file_name = DOWNLOAD_DIR + '/' + target_url.split('/')[-1]
    with open(file_name, 'wb') as f:
        f.write(file_binary.content)
    print('Downloaded:', file_name)

# Latest version
download_latest_version(_versions[0], 'chrome')
download_latest_version(_versions[0], 'chromedriver')

print('Input sudo password')
sudo_password = (getpass.getpass() + '\n').encode()

def _install_chromedriver(sudo_password):
    print('Installing chromedriver...', flush=True)
    _platform = get_current_machine_platform()
    _zip_file_path = f'{DOWNLOAD_DIR}/chromedriver-{_platform}.zip'
    _extracted_dir_path = f'{DOWNLOAD_DIR}/chromedriver-{_platform}'

    with set_directory(DOWNLOAD_DIR):
        if os.path.exists(_extracted_dir_path):
            shutil.rmtree(_extracted_dir_path, ignore_errors=True)
        unzip_output = subprocess.run(
            ['unzip', _zip_file_path],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT
        ).stdout.decode('utf-8').strip()
        print(unzip_output)

        exists_chromedriver_path = shutil.which('chromedriver')
        print('exists_chromedriver_path:', exists_chromedriver_path)

        install_output = subprocess.run(
            ['sudo', '-S', 'mv', f'{_extracted_dir_path}/chromedriver',
             exists_chromedriver_path], input=sudo_password, check=True
        )
        print('Chromedriver installed.', install_output)

_install_chromedriver(sudo_password)

def _install_chrome(sudo_password):
    print('Installing chrome...', flush=True)
    _platform = get_current_machine_platform()
    _zip_file_path = f'{DOWNLOAD_DIR}/chrome-{_platform}.zip'
    _extracted_dir_path = f'{DOWNLOAD_DIR}/chrome-{_platform}'
    _installed_chrome_path = '/Applications/Google Chrome for Testing.app'

    with set_directory(DOWNLOAD_DIR):
        if os.path.exists(_extracted_dir_path):
            shutil.rmtree(_extracted_dir_path, ignore_errors=True)

        if os.path.exists(_installed_chrome_path):
            shutil.rmtree(_installed_chrome_path, ignore_errors=True)

        unzip_output = subprocess.run(
            ['unzip', _zip_file_path],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT
        ).stdout.decode('utf-8').strip()
        print(unzip_output)

        install_output = subprocess.run(
            ['sudo', '-S', 'mv',
             f'chrome-{_platform}/Google Chrome for Testing.app',
             '/Applications/'], input=sudo_password, check=True
        )

        print('Chrome for Testing installed.', install_output)

_install_chrome(sudo_password)