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, logger=None):
 31        super(LocalFileIO, self).__init__()
 32        self.logger = logger
 33        self.mode = mode
 34        # Convert remote path format to local operating system path format 
 35        self.path = path.replace(ntpath.sep, os.path.sep)
 36        self.dir = None
 37        self.debug = False
 38        self.expected_size = expected_size
 39        self.keepRemotePath = keepRemotePath
 40
 41        # Write to local (read remote)
 42        if self.mode in ["wb"]:
 43            self.dir = '.' + os.path.sep
 44            if keepRemotePath:
 45                self.dir += os.path.dirname(self.path)
 46
 47            if not os.path.exists(self.dir):
 48                self.logger.debug("[debug] Creating local directory '%s'" % self.dir)
 49                os.makedirs(self.dir)
 50
 51            self.logger.debug("[debug] Openning local '%s' with mode '%s'" % (self.path, self.mode))
 52            
 53            try:
 54                self.fd = open(self.dir + os.path.sep + os.path.basename(self.path), self.mode)
 55            except PermissionError as err:
 56                self.fd = None
 57
 58        # Write to remote (read local)
 59        elif self.mode in ["rb"]:
 60            if ntpath.sep in self.path:
 61                self.dir = os.path.dirname(self.path)
 62
 63            self.logger.debug("[debug] Openning local '%s' with mode '%s'" % (self.path, self.mode))
 64            
 65            try:
 66                self.fd = open(self.path, self.mode)
 67            except PermissionError as err:
 68                self.fd = None
 69
 70            if self.fd is not None:
 71                if self.expected_size is None:
 72                    self.expected_size = os.path.getsize(filename=self.path)
 73
 74        # Create progress bar
 75        if self.expected_size is not None:
 76            self.__progress = Progress(
 77                TextColumn("[bold blue]{task.description}", justify="right"),
 78                BarColumn(bar_width=None),
 79                "[progress.percentage]{task.percentage:>3.1f}%",
 80                "•",
 81                DownloadColumn(),
 82                "•",
 83                TransferSpeedColumn(),
 84                "•",
 85                TimeRemainingColumn(),
 86            )
 87            self.__progress.start()
 88            self.__task = self.__progress.add_task(
 89                description="'%s'" % os.path.basename(self.path),
 90                start=True,
 91                total=self.expected_size,
 92                visible=True
 93            )
 94
 95    def write(self, data):
 96        """
 97        Writes data to the file.
 98
 99        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.
100
101        Args:
102            data (bytes): The data to be written to the file.
103
104        Returns:
105            int: The number of bytes written.
106        """
107
108        if self.fd is not None:
109            if self.expected_size is not None:
110                self.__progress.update(self.__task, advance=len(data))
111            return self.fd.write(data)
112        else:
113            return 0
114    
115    def read(self, size):
116        """
117        Reads a specified amount of data from the file.
118
119        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.
120
121        Args:
122            size (int): The number of bytes to read from the file.
123
124        Returns:
125            bytes: The data read from the file.
126        """
127
128        if self.fd is not None:
129            read_data = self.fd.read(size)
130            if self.expected_size is not None:
131                self.__progress.update(self.__task, advance=len(read_data))
132            return read_data
133        else:
134            return b""
135
136    def close(self, remove=False):
137        """
138        Closes the file descriptor and optionally removes the file.
139
140        This method ensures that the file descriptor is properly closed and the file is removed if specified.
141        It also stops the progress bar if it was initiated and cleans up the object by deleting it.
142
143        Args:
144            remove (bool): If True, the file at the path will be removed after closing the file descriptor.
145        """
146
147        if self.fd is not None:
148            self.fd.close()
149
150        if remove:
151            try:
152                os.remove(path=self.path)
153            except (PermissionError, FileNotFoundError) as err:
154                pass
155
156        if self.expected_size is not None:
157            self.__progress.stop()
158        
159        del self
160
161    def set_error(self, message):
162        """
163        Sets an error message in the progress bar's description and modifies the progress bar to show only essential columns.
164
165        This method is used to communicate error states or important messages directly in the progress bar interface.
166        It updates the task description with the provided message and simplifies the progress bar to show only the text
167        and download columns, removing other elements like speed and time remaining which may not be relevant in an error state.
168
169        Args:
170            message (str): The error or status message to display in the progress bar.
171        """
172
173        self.__progress.tasks[0].description = message
174        self.__progress.columns = [
175            TextColumn("[bold blue]{task.description}", justify="right"),
176            BarColumn(bar_width=None),
177            "•",
178            DownloadColumn(),
179        ]
180        self.__progress.update(self.__task, advance=0)