Download file given Google Drive URL

gdrive_download.py

"""
Download file from Google Drive URL
"""
import os
import time

import click
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC


@click.command()
@click.option("--url")
@click.option("--download_dir", default=".")
@click.option("--timeout", default=300)
def download(url, download_dir, timeout):
    firefox_options = FirefoxOptions()
    firefox_options.add_argument("--headless")
    firefox_options.set_preference("browser.download.folderList", 2)
    firefox_options.set_preference("browser.download.manager.showWhenStarting", False)
    firefox_options.set_preference(
        "browser.download.dir", os.path.abspath(download_dir)
    )
    firefox_options.set_preference(
        "browser.helperApps.neverAsk.saveToDisk", "application/octet-stream"
    )
    driver = webdriver.Firefox(options=firefox_options)

    try:
        driver.get(url)
        WebDriverWait(driver, 10).until(
            EC.element_to_be_clickable((By.ID, "uc-download-link"))
        ).click()
        if wait_for_download_complete(download_dir, timeout):
            print("Download completed successfully.")
        else:
            print("Download did not complete within the expected time.")
    finally:
        driver.quit()

        
def wait_for_download_complete(download_dir, timeout):
    start_time = time.time()
    while True:
        if not any(f.endswith(".part") for f in os.listdir(download_dir)):
            return True
        elif time.time() - start_time > timeout:
            return False
        else:
            time.sleep(1)


if __name__ == "__main__":
    download()