#! /usr/bin/python

import urllib.request
import urllib.parse
import re
import os
import subprocess
import ssl
import sys
import argparse
import tempfile
import textwrap

try:
  from tqdm import tqdm
  TQDM_AVAILABLE = True
except ImportError:
  TQDM_AVAILABLE = False

class DownloadProgressBar:
  def __init__(self, total=0, unit='iB', unit_scale=True, unit_divisor=1024):
    self.pbar = None
    self.total = total
    self.unit = unit
    self.unit_scale = unit_scale
    self.unit_divisor = unit_divisor

  def __call__(self, downloaded, block_size, total_size):
    if not self.pbar:
      self.total = total_size
      self.pbar = tqdm(total=total_size, unit=self.unit, unit_scale=self.unit_scale, unit_divisor=self.unit_divisor)
    
    self.pbar.update(downloaded - self.pbar.n)

    if downloaded >= total_size:
      self.pbar.close()

def simple_progress_bar(downloaded, block_size, total_size):
  progress = min(round(downloaded / total_size * 100, 2), 100)
  sys.stdout.write(f"\rDownload progress: {progress:6.2f}%")
  sys.stdout.flush()

def create_ssl_context():
  cert_chain = textwrap.dedent("""
      -----BEGIN CERTIFICATE-----
      MIIFVzCCAz+gAwIBAgINAgPlk28xsBNJiGuiFzANBgkqhkiG9w0BAQwFADBHMQsw
      CQYDVQQGEwJVUzEiMCAGA1UEChMZR29vZ2xlIFRydXN0IFNlcnZpY2VzIExMQzEU
      MBIGA1UEAxMLR1RTIFJvb3QgUjEwHhcNMTYwNjIyMDAwMDAwWhcNMzYwNjIyMDAw
      MDAwWjBHMQswCQYDVQQGEwJVUzEiMCAGA1UEChMZR29vZ2xlIFRydXN0IFNlcnZp
      Y2VzIExMQzEUMBIGA1UEAxMLR1RTIFJvb3QgUjEwggIiMA0GCSqGSIb3DQEBAQUA
      A4ICDwAwggIKAoICAQC2EQKLHuOhd5s73L+UPreVp0A8of2C+X0yBoJx9vaMf/vo
      27xqLpeXo4xL+Sv2sfnOhB2x+cWX3u+58qPpvBKJXqeqUqv4IyfLpLGcY9vXmX7w
      Cl7raKb0xlpHDU0QM+NOsROjyBhsS+z8CZDfnWQpJSMHobTSPS5g4M/SCYe7zUjw
      TcLCeoiKu7rPWRnWr4+wB7CeMfGCwcDfLqZtbBkOtdh+JhpFAz2weaSUKK0Pfybl
      qAj+lug8aJRT7oM6iCsVlgmy4HqMLnXWnOunVmSPlk9orj2XwoSPwLxAwAtcvfaH
      szVsrBhQf4TgTM2S0yDpM7xSma8ytSmzJSq0SPly4cpk9+aCEI3oncKKiPo4Zor8
      Y/kB+Xj9e1x3+naH+uzfsQ55lVe0vSbv1gHR6xYKu44LtcXFilWr06zqkUspzBmk
      MiVOKvFlRNACzqrOSbTqn3yDsEB750Orp2yjj32JgfpMpf/VjsPOS+C12LOORc92
      wO1AK/1TD7Cn1TsNsYqiA94xrcx36m97PtbfkSIS5r762DL8EGMUUXLeXdYWk70p
      aDPvOmbsB4om3xPXV2V4J95eSRQAogB/mqghtqmxlbCluQ0WEdrHbEg8QOB+DVrN
      VjzRlwW5y0vtOUucxD/SVRNuJLDWcfr0wbrM7Rv1/oFB2ACYPTrIrnqYNxgFlQID
      AQABo0IwQDAOBgNVHQ8BAf8EBAMCAYYwDwYDVR0TAQH/BAUwAwEB/zAdBgNVHQ4E
      FgQU5K8rJnEaK0gnhS9SZizv8IkTcT4wDQYJKoZIhvcNAQEMBQADggIBAJ+qQibb
      C5u+/x6Wki4+omVKapi6Ist9wTrYggoGxval3sBOh2Z5ofmmWJyq+bXmYOfg6LEe
      QkEzCzc9zolwFcq1JKjPa7XSQCGYzyI0zzvFIoTgxQ6KfF2I5DUkzps+GlQebtuy
      h6f88/qBVRRiClmpIgUxPoLW7ttXNLwzldMXG+gnoot7TiYaelpkttGsN/H9oPM4
      7HLwEXWdyzRSjeZ2axfG34arJ45JK3VmgRAhpuo+9K4l/3wV3s6MJT/KYnAK9y8J
      ZgfIPxz88NtFMN9iiMG1D53Dn0reWVlHxYciNuaCp+0KueIHoI17eko8cdLiA6Ef
      MgfdG+RCzgwARWGAtQsgWSl4vflVy2PFPEz0tv/bal8xa5meLMFrUKTX5hgUvYU/
      Z6tGn6D/Qqc6f1zLXbBwHSs09dR2CQzreExZBfMzQsNhFRAbd03OIozUhfJFfbdT
      6u9AWpQKXCBfTkBdYiJ23//OYb2MI3jSNwLgjt7RETeJ9r/tSQdirpLsQBqvFAnZ
      0E6yove+7u7Y/9waLd64NnHi/Hm3lCXRSHNboTXns5lndcEZOitHTtNCjv0xyBZm
      2tIMPNuzjsmhDYAPexZ3FL//2wmUspO8IFgV6dtxQ/PeEMMA3KgqlbbC1j+Qa3bb
      bP6MvPJwNQzcmRk13NfIRmPVNnGuV/u3gm3c
      -----END CERTIFICATE-----
  """).strip()

  context = ssl.create_default_context(cadata=cert_chain)

  return context

def download_large_file(file_id, output_filename, ssl_context):
  url = f"https://drive.google.com/uc?export=download&id={file_id}"
  req = urllib.request.Request(url, method="GET")
  
  try:
      with urllib.request.urlopen(req, context=ssl_context) as response:
          response_content = response.read().decode('utf-8')
  except urllib.error.URLError as e:
      print(f"Error accessing Google Drive: {e}")
      return False

  form_action = re.search(r'<form id="download-form" action="([^"]*)"', response_content)
  confirm = re.search(r'name="confirm" value="([^"]*)"', response_content)
  uuid_value = re.search(r'name="uuid" value="([^"]*)"', response_content)

  if not all([form_action, confirm, uuid_value]):
      print("Failed to extract form details. The download process might have changed.")
      return False

  download_url = (f"{form_action.group(1)}?"
                  f"id={file_id}&"
                  f"export=download&"
                  f"confirm={confirm.group(1)}&"
                  f"uuid={uuid_value.group(1)}")

  print(f"Downloading file to {output_filename}...")
  try:
    with urllib.request.urlopen(download_url, context=ssl_context) as response, open(output_filename, 'wb') as out_file:
      file_size = int(response.headers.get('Content-Length', 0))
      print(file_size)
      block_size = 8192
      
      if TQDM_AVAILABLE:
        progress_bar = tqdm(total=file_size, unit='iB', unit_scale=True)
      else:
        progress_bar = simple_progress_bar
      
      size = 0
      while True:
        buffer = response.read(block_size)
        if not buffer:
          break
        size += len(buffer)
        out_file.write(buffer)
        if TQDM_AVAILABLE:
          progress_bar.update(len(buffer))
        else:
          progress_bar(size, block_size, file_size)
      
      if TQDM_AVAILABLE:
        progress_bar.close()
      else:
        print()  # New line after progress bar
    
    print("Download completed.")
    return True
  except urllib.error.URLError as e:
    print(f"\nError downloading file: {e}")
    return False

def extract_file(filename, extract_path):
  print(f"Extracting \033[94m{filename}\033[0m to \033[96m{extract_path}\033[0m...")
  try:
    subprocess.run(["7z", "x", filename, f"-o{extract_path}"], check=True)
    print("Extraction completed.")
    return True
  except subprocess.CalledProcessError:
    print("Extraction failed. Make sure 7z is installed and in your PATH.")
    return False

def check_writable(path):
  try:
    testfile = tempfile.TemporaryFile(dir=path)
    testfile.close()
  except OSError:
    return False
  return True

def print_usage():
  print("Usage: python3 download_script.py [--datadir PATH]")
  print("Options:")
  print("  --datadir PATH, -d PATH    Specify the path to the Stealth data directory")
  print("                             (default: ~/.StealthCoin)")

def main():
  parser = argparse.ArgumentParser(description="Download and extract Stealth bootstrap file.")
  parser.add_argument('--datadir', '-d', type=str, help='Path to the Stealth data directory')
  args = parser.parse_args()

  if not check_writable('.'):
    print("Error: Current working directory is not writable.")
    sys.exit(1)

  if args.datadir:
    data_dir = os.path.abspath(args.datadir)
  else:
    data_dir = os.path.expanduser("~/.StealthCoin")

  if not os.path.exists(data_dir):
    print(f"Error: Cannot find the Stealth data directory at {data_dir}")
    print_usage()
    sys.exit(1)

  file_id = "1WNmH1DiO4-_D-UYupoyFXM-rKMlGVTet"
  output_filename = "xst-bootstrap.7z"

  print("This script will download a large file (3.0G) from Google Drive.")
  print("Google Drive can't scan this file for viruses due to its size.")
  
  ssl_context = create_ssl_context()
  print("Using bundled Google Drive certificate chain for secure connection.")
  
  user_input = input("\033[1mDo you want to proceed with the download? (y/n):\033[0m ").lower()

  if user_input == 'y':
    if download_large_file(file_id, output_filename, ssl_context):
      user_input = input("\033[1mDo you want to extract the downloaded file? (y/n):\033[0m ").lower()
      if user_input == 'y':
        extract_file(output_filename, data_dir)
  else:
    print("\033[31mDownload cancelled.\033[0m")

if __name__ == "__main__":
  main()
