src/utils.py

#!/usr/bin/env python
# pylint: disable=C0116,W0613
# This program is dedicated to the public domain under the CC0 license.

import json
import os
from os.path import exists
from typing import Optional
import urllib.request
import socket
import pickle
import subprocess
import asyncio
import psutil
import time
from datetime import datetime
from telegram import Message
from telegram.error import TelegramError, RetryAfter
from telegram.ext import CallbackContext

class ProgressTracker:
    def __init__(self, loop: Optional[asyncio.AbstractEventLoop], context: CallbackContext, msg: Message, filename: str, update_interval: float = 2.0):
        self.msg = msg
        self.context = context
        self.filename = filename
        self.update_interval = update_interval
        self.last_update = 0.0
        self.total = 0.0
        self.loop = asyncio.get_event_loop() if loop == None else loop

    async def progress(self, current, total):
        self.total = total

        now = self.loop.time()
        if now - self.last_update < self.update_interval:
            return

        self.last_update = now

        try:
            _, bar, perc = print_progress_bar(current, total, decimals=0, length=50)
            text = (
                f"📥 Downloading...\n{self.filename}\n"
                f"{perc}% |{bar}|\n"
                f"{current//(1024)} KB / {total//(1024)} KB"
            )
            await self.context.bot.edit_message_text(text=text, message_id=self.msg.message_id)
            
        except RetryAfter as e:
            await asyncio.sleep(float(e.value)) # type: ignore
        except TelegramError:
            pass
            
    async def finished(self):

        try:
            _, bar, perc = print_progress_bar(self.total, self.total, decimals=0, length=50)
            text = (
                f"📥 Downloading...\n{self.filename}\n"
                f"{perc}% |{bar}|\n"
                f"{self.total//(1024)} KB / {self.total//(1024)} KB"
            )
            await self.context.bot.edit_message_text(text=text, message_id=self.msg.message_id)
            
        except RetryAfter as e:
            await asyncio.sleep(float(e.value)) # type: ignore
        except TelegramError:
            pass


def save_obj(obj, name):
    filename = os.path.dirname(os.path.realpath(
        __file__)) + "/obj/" + name + '.pkl'
    with open(filename, 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


def load_obj(name):
    filename = os.path.dirname(os.path.realpath(
        __file__)) + "/obj/" + name + '.pkl'
    try:
        with open(filename, 'rb') as f:
            return pickle.load(f)
    except FileNotFoundError:
        return None


def load_json(filename):
    try:
        with open(filename, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        print("load_json: file not found: " + filename)
        return dict()


def toJSON(obj):
    return json.dumps(obj)


def run_command(cmd, output=False, timeout=None):

    print(f"cmd: {cmd}, output: {output}, timeout: {timeout}")
    outputStr = "Done"

    process = subprocess.Popen(
        [cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        out, _ = process.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        print("timeout!!")
        # process.kill()
        parent = psutil.Process(process.pid)
        for child in parent.children(recursive=True):
            child.kill()
        parent.kill()
        out, _ = process.communicate()
    if output:
        outputStr = out.decode('UTF-8')
    return outputStr

async def run_command_en(cmd, message, timeout=None):

    if timeout == None:
        timeout = 5*60

    print(f"Stream output cmd: {cmd}, timeout: {timeout}")

    # Set the timeout duration (5 minutes)
    timeout_duration = timeout  # 5 minutes in seconds

    # Record the start time
    start_time = time.time()

    # Run the shell script
    process = subprocess.Popen(
        [cmd], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, bufsize=1)

    # Continuously read and print the output of the inner script
    try:

        while True:
            # Calculate elapsed time
            elapsed_time = time.time() - start_time

            # Check if the timeout duration has been exceeded
            if elapsed_time > timeout_duration:
                print("Timeout exceeded. Terminating the inner script.")
                break
            
            output = ''
            if process.stdout:
                process.stdout.flush()
                output = process.stdout.readline()

            print(f"output cmd: {output}")
            if output == '' and process.poll() is not None:
                break
            if output.strip():
                await message.reply_text(f"{output}")
    except KeyboardInterrupt:
        print("Outer script interrupted by user.")

    # finally:
    #     # Ensure the process is terminated if not already
    #     if process.poll() is None:
    #         process.terminate()
    #         process.wait()

    return None


def get_public_ip():
    ip = ""
    try:
        # data = json.loads(urllib.request.urlopen("https://api.ipify.org").read().decode('utf8'))
        ip = urllib.request.urlopen(
            "https://api.ipify.org").read().decode('utf8')
    except:
        ip = "err"

    return ip


def get_local_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        # doesn't even have to be reachable
        s.connect(('10.255.255.255', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP


def fstr(template):
    local_ip = get_local_ip()

    return eval(f"f\"{template}\"")


# Print iterations progress
def print_progress_bar(current, total, prefix='', suffix='', decimals=1, length=100, fill='#', notFill='.', printEnd="\r"):
    """
    Call in a loop to create terminal progress bar
    @params:
        current   - Required  : current current (Int)
        total       - Required  : total size (Int)
        prefix      - Optional  : prefix string (Str)
        suffix      - Optional  : suffix string (Str)
        decimals    - Optional  : positive number of decimals in percent complete (Int)
        length      - Optional  : character length of bar (Int)
        fill        - Optional  : bar fill character (Str)
        printEnd    - Optional  : end character (e.g. "\r", "\r\n") (Str)
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 *
                                                     (current / float(total)))
    filledLength = int(length * current // total)
    bar = fill * filledLength + notFill * (length - filledLength)
    # progress_string = f'\r{prefix} |{bar}| {percent}% {suffix}'
    progress_string = f'\r{prefix} {percent}%|{bar}| {current}/{total} {suffix}'
    # print(progress_string)
    # Print New Line on Complete
    # if current == total:
    #     print()
    # else:

    return progress_string, bar, percent


# Function to generate a unique filename with timestamp
def get_unique_filepath(base_path):
    directory = os.path.dirname(base_path)
    filename = os.path.basename(base_path)
    name, ext = os.path.splitext(filename)

    counter = 0
    while os.path.exists(base_path):
        counter += 1
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        base_path = os.path.join(directory, f"{name}_{timestamp}{ext}")

    return base_path