helper to download v2ray daily updated rules

update_geodb.py

#-*- coding: UTF-8 -*-
import os
import sys
import time
import logging
import socket
import urllib.request
import base64
import hashlib
import tempfile

from urllib.parse import urlparse
from pathlib import Path
from functools import partial

import shutil
import requests
from tqdm import tqdm

GEODB_LATEST = 'https://github.com/Loyalsoldier/v2ray-rules-dat/releases/latest'
GEODB_URL = 'https://github.com/Loyalsoldier/v2ray-rules-dat/releases/download/'
DB_LIST = ['geoip.dat', 'geosite.dat']
HEADERS = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36'}
PROXIES = {
   'http': 'http://127.0.0.1:10809',
   'https': 'http://127.0.0.1:10809',
}

BINGWEB = "http://www.bing.com"
CONNECTION_TRIES = 120
LOGFILE = "geodb.log"
CHUNK_SIZE = 1024 * 64

## for v2ray versions before 5.39
# GEODB_SAVE_PATH = Path("D:\\bin\\v2rayN")

## for v2ray versions 6.0+
GEODB_SAVE_PATH = Path("D:\\bin\\v2rayN\\bin")

logger = logging.getLogger(__name__)

def file_digest(file):
    d = hashlib.md5()
    for buf in iter(partial(file.read, 4096), b''):
        d.update(buf)
    return d.digest()

def setup_logging(file_log=False):
    formatter = logging.Formatter('%(name)-12s %(asctime)s %(levelname)-8s %(message)s', '%a, %d %b %Y %H:%M:%S')
    if file_log:
        logfile = Path(__file__).parent / LOGFILE
        file_handler = logging.FileHandler(logfile)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    stream_handler = logging.StreamHandler(sys.stderr)
    logger.addHandler(stream_handler)
    logger.setLevel(logging.INFO)

def is_internet_ready():
    logger.info("Probing internet availability...")
    try:
        urllib.request.urlopen(BINGWEB, timeout=1)
    except:
        return False
    else:
        logger.info("Internet is ready")
        return True

def is_proxy_ready():
    proxy = urlparse(PROXIES['http'])
    host = proxy.hostname
    port = proxy.port
    logger.info(f"Probing proxy: {proxy.netloc}...")

    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.connect((host, port))
    except:
        return False
    else:
        logger.info("Proxy is ready")
        return True

def get_latest_tag():
    try:
        r = requests.get(GEODB_LATEST, proxies=PROXIES)
    except (requests.exceptions.ConnectTimeout, requests.exceptions.ConnectionError) as msg:
        logger.exception(msg)
        return None
    logger.info(f"{r.url}")
    o = urlparse(r.url)
    return o.path.rsplit('/', 1)[-1]

def main():
    logger.info("")
    tag = get_latest_tag()
    if not tag:
        return 1

    for f in DB_LIST:
        logger.info("")
        file_to_save = GEODB_SAVE_PATH / f

        r = requests.get(GEODB_URL + tag + '/' + f, proxies=PROXIES, headers=HEADERS, stream=True)

        length = int(r.headers['Content-Length'])
        if not length:
            logger.error("file length 0")
            return 1
        logger.info(f"File size: {length} bytes")

        m = r.headers['Content-MD5']
        logger.debug(f"File md5: {m}")

        if file_to_save.exists() and file_to_save.stat().st_size == length:
            digest = None
            with open(file_to_save, 'rb') as f:
                try:
                    digest = hashlib.file_digest(f, 'md5').digest()
                except AttributeError:  # file_digest is introduced since py3.11
                    digest = file_digest(f)
            if digest == base64.b64decode(m):
                logger.info("Local file is up to date")
                continue
            logger.debug("Local file was out of date")

        with tempfile.TemporaryDirectory() as dtemp:
            temp_file = Path(dtemp.name) / f
            with open(temp_file, 'wb') as db:
                with tqdm(total=length, unit='byte', unit_scale=True) as pbar:
                    for chunk in r.iter_content(chunk_size=CHUNK_SIZE):
                        db.write(chunk)
                        pbar.update(CHUNK_SIZE)

                    if temp_file.stat().sz_size == length:
                        shutil.copy2(temp_file, file_to_save)
                        logger.info(f'Copied {temp_file} to {file_to_save}')
                    else:
                        logger.error('Failed to download the file')
                        return 1

    return 0

if __name__ == "__main__":
    setup_logging()

    for cnt in range(CONNECTION_TRIES):
        if is_internet_ready():
            break
        time.sleep(2)
        if cnt % 5 == 0:
            logger.error("Internet connection is not ready, retrying...")
    else:
        logger.error("Internet connection is not ready")
        sys.exit(1)

    for cnt in range(CONNECTION_TRIES):
        if is_proxy_ready():
            break
        time.sleep(2)
        if cnt % 5 == 0:
            logger.error("proxy seems not ready, retrying...")
    else:
        logger.error("proxy seems not ready")
        sys.exit(1)

    sys.exit(main())