smbclientng.core.LocalFileIO

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3# File name          : LocalFileIO.py
  4# Author             : Podalirius (@podalirius_)
  5# Date created       : 23 may 2024
  6
  7
  8import os
  9import ntpath
 10from rich.progress import BarColumn, DownloadColumn, Progress, TextColumn, TimeRemainingColumn, TransferSpeedColumn
 11
 12
 13class LocalFileIO(object):
 14    """
 15    Class LocalFileIO is designed to handle local file input/output operations within the smbclient-ng tool.
 16    It provides functionalities to open, read, write, and manage progress of file operations based on the expected size of the file.
 17
 18    Attributes:
 19        mode (str): The mode in which the file should be opened (e.g., 'rb', 'wb').
 20        path (str): The path to the file that needs to be handled.
 21        expected_size (int, optional): The expected size of the file in bytes. This is used to display progress.
 22        debug (bool): Flag to enable debug mode which provides additional output during operations.
 23
 24    Methods:
 25        __init__(self, mode, path=None, expected_size=None, debug=False): Initializes the LocalFileIO instance.
 26        write(self, data): Writes data to the file and updates the progress bar if expected size is provided.
 27        read(self, size): Reads data from the file up to the specified size and updates the progress bar if expected size is provided.
 28    """
 29
 30    def __init__(self, mode, path=None, expected_size=None, keepRemotePath=False, debug=False):
 31        super(LocalFileIO, self).__init__()
 32        self.mode = mode
 33        # Convert remote path format to local operating system path format 
 34        self.path = path.replace(ntpath.sep, os.path.sep)
 35        self.dir = None
 36        self.debug = False
 37        self.expected_size = expected_size
 38        self.keepRemotePath = keepRemotePath
 39
 40        # Write to local (read remote)
 41        if self.mode in ["wb"]:
 42            self.dir = '.' + os.path.sep
 43            if keepRemotePath:
 44                self.dir += os.path.dirname(self.path)
 45
 46            if not os.path.exists(self.dir):
 47                if self.debug:
 48                    print("[debug] Creating local directory '%s'" % self.dir)
 49                os.makedirs(self.dir)
 50
 51            if self.debug:
 52                print("[debug] Openning local '%s' with mode '%s'" % (self.path, self.mode))
 53            
 54            self.fd = open(self.dir + os.path.sep + os.path.basename(self.path), self.mode)
 55
 56        # Write to remote (read local)
 57        elif self.mode in ["rb"]:
 58            if ntpath.sep in self.path:
 59                self.dir = os.path.dirname(self.path)
 60
 61            if self.debug:
 62                print("[debug] Openning local '%s' with mode '%s'" % (self.path, self.mode))
 63            self.fd = open(self.path, self.mode)
 64
 65            if self.expected_size is None:
 66                self.expected_size = os.path.getsize(filename=self.path)
 67
 68        # Create progress bar
 69        if self.expected_size is not None:
 70            self.__progress = Progress(
 71                TextColumn("[bold blue]{task.description}", justify="right"),
 72                BarColumn(bar_width=None),
 73                "[progress.percentage]{task.percentage:>3.1f}%",
 74                "•",
 75                DownloadColumn(),
 76                "•",
 77                TransferSpeedColumn(),
 78                "•",
 79                TimeRemainingColumn(),
 80            )
 81            self.__progress.start()
 82            self.__task = self.__progress.add_task(
 83                description="'%s'" % os.path.basename(self.path),
 84                start=True,
 85                total=self.expected_size,
 86                visible=True
 87            )
 88
 89    def write(self, data):
 90        """
 91        Writes data to the file.
 92
 93        This method writes the specified data to the file and updates the progress bar with the amount of data written if the expected size is set.
 94
 95        Args:
 96            data (bytes): The data to be written to the file.
 97
 98        Returns:
 99            int: The number of bytes written.
100        """
101
102        if self.expected_size is not None:
103            self.__progress.update(self.__task, advance=len(data))
104        return self.fd.write(data)
105    
106    def read(self, size):
107        """
108        Reads a specified amount of data from the file.
109
110        This method reads data from the file based on the size specified. It also updates the progress bar with the amount of data read if the expected size is set.
111
112        Args:
113            size (int): The number of bytes to read from the file.
114
115        Returns:
116            bytes: The data read from the file.
117        """
118
119        read_data = self.fd.read(size)
120        if self.expected_size is not None:
121            self.__progress.update(self.__task, advance=len(read_data))
122        return read_data
123
124    def close(self, remove=False):
125        """
126        Closes the file descriptor and optionally removes the file.
127
128        This method ensures that the file descriptor is properly closed and the file is removed if specified.
129        It also stops the progress bar if it was initiated and cleans up the object by deleting it.
130
131        Args:
132            remove (bool): If True, the file at the path will be removed after closing the file descriptor.
133        """
134
135        self.fd.close()
136
137        if remove:
138            os.remove(path=self.path)
139
140        if self.expected_size is not None:
141            self.__progress.stop()
142        
143        del self
144
145    def set_error(self, message):
146        """
147        Sets an error message in the progress bar's description and modifies the progress bar to show only essential columns.
148
149        This method is used to communicate error states or important messages directly in the progress bar interface.
150        It updates the task description with the provided message and simplifies the progress bar to show only the text
151        and download columns, removing other elements like speed and time remaining which may not be relevant in an error state.
152
153        Args:
154            message (str): The error or status message to display in the progress bar.
155        """
156
157        self.__progress.tasks[0].description = message
158        self.__progress.columns = [
159            TextColumn("[bold blue]{task.description}", justify="right"),
160            BarColumn(bar_width=None),
161            "•",
162            DownloadColumn(),
163        ]
164        self.__progress.update(self.__task, advance=0)