codekingpro/portable-devtools

1"""2This module contains SMBFileSystem class responsible for handling access to3Windows Samba network shares by using package smbprotocol4"""5 6import datetime7import re8import uuid9from stat import S_ISDIR, S_ISLNK10 11import smbclient12import smbprotocol.exceptions13 14from .. import AbstractFileSystem15from ..utils import infer_storage_options16 17# ! pylint: disable=bad-continuation18 19 20class SMBFileSystem(AbstractFileSystem):21    """Allow reading and writing to Windows and Samba network shares.22 23    When using `fsspec.open()` for getting a file-like object the URI24    should be specified as this format:25    ``smb://workgroup;user:password@server:port/share/folder/file.csv``.26 27    Example::28 29        >>> import fsspec30        >>> with fsspec.open(31        ...     'smb://myuser:mypassword@myserver.com/' 'share/folder/file.csv'32        ... ) as smbfile:33        ...     df = pd.read_csv(smbfile, sep='|', header=None)34 35    Note that you need to pass in a valid hostname or IP address for the host36    component of the URL. Do not use the Windows/NetBIOS machine name for the37    host component.38 39    The first component of the path in the URL points to the name of the shared40    folder. Subsequent path components will point to the directory/folder/file.41 42    The URL components ``workgroup`` , ``user``, ``password`` and ``port`` may be43    optional.44 45    .. note::46 47        For working this source require `smbprotocol`_ to be installed, e.g.::48 49            $ pip install smbprotocol50            # or51            # pip install smbprotocol[kerberos]52 53    .. _smbprotocol: https://github.com/jborean93/smbprotocol#requirements54 55    Note: if using this with the ``open`` or ``open_files``, with full URLs,56    there is no way to tell if a path is relative, so all paths are assumed57    to be absolute.58    """59 60    protocol = "smb"61 62    # pylint: disable=too-many-arguments63    def __init__(64        self,65        host,66        port=None,67        username=None,68        password=None,69        timeout=60,70        encrypt=None,71        share_access=None,72        register_session_retries=4,73        register_session_retry_wait=1,74        register_session_retry_factor=10,75        auto_mkdir=False,76        **kwargs,77    ):78        """79        You can use _get_kwargs_from_urls to get some kwargs from80        a reasonable SMB url.81 82        Authentication will be anonymous or integrated if username/password are not83        given.84 85        Parameters86        ----------87        host: str88            The remote server name/ip to connect to89        port: int or None90            Port to connect with. Usually 445, sometimes 139.91        username: str or None92            Username to connect with. Required if Kerberos auth is not being used.93        password: str or None94            User's password on the server, if using username95        timeout: int96            Connection timeout in seconds97        encrypt: bool98            Whether to force encryption or not, once this has been set to True99            the session cannot be changed back to False.100        share_access: str or None101            Specifies the default access applied to file open operations102            performed with this file system object.103            This affects whether other processes can concurrently open a handle104            to the same file.105 106            - None (the default): exclusively locks the file until closed.107            - 'r': Allow other handles to be opened with read access.108            - 'w': Allow other handles to be opened with write access.109            - 'd': Allow other handles to be opened with delete access.110        register_session_retries: int111            Number of retries to register a session with the server. Retries are not performed112            for authentication errors, as they are considered as invalid credentials and not network113            issues. If set to negative value, no register attempts will be performed.114        register_session_retry_wait: int115            Time in seconds to wait between each retry. Number must be non-negative.116        register_session_retry_factor: int117            Base factor for the wait time between each retry. The wait time118            is calculated using exponential function. For factor=1 all wait times119            will be equal to `register_session_retry_wait`. For any number of retries,120            the last wait time will be equal to `register_session_retry_wait` and for retries>1121            the first wait time will be equal to `register_session_retry_wait / factor`.122            Number must be equal to or greater than 1. Optimal factor is 10.123        auto_mkdir: bool124            Whether, when opening a file, the directory containing it should125            be created (if it doesn't already exist). This is assumed by pyarrow126            and zarr-python code.127        """128        super().__init__(**kwargs)129        self.host = host130        self.port = port131        self.username = username132        self.password = password133        self.timeout = timeout134        self.encrypt = encrypt135        self.temppath = kwargs.pop("temppath", "")136        self.share_access = share_access137        self.register_session_retries = register_session_retries138        if register_session_retry_wait < 0:139            raise ValueError(140                "register_session_retry_wait must be a non-negative integer"141            )142        self.register_session_retry_wait = register_session_retry_wait143        if register_session_retry_factor < 1:144            raise ValueError(145                "register_session_retry_factor must be a positive "146                "integer equal to or greater than 1"147            )148        self.register_session_retry_factor = register_session_retry_factor149        self.auto_mkdir = auto_mkdir150        self._connect()151 152    @property153    def _port(self):154        return 445 if self.port is None else self.port155 156    def _connect(self):157        import time158 159        if self.register_session_retries <= -1:160            return161 162        retried_errors = []163 164        wait_time = self.register_session_retry_wait165        n_waits = (166            self.register_session_retries - 1167        )  # -1 = No wait time after the last retry168        factor = self.register_session_retry_factor169 170        # Generate wait times for each retry attempt.171        # Wait times are calculated using exponential function. For factor=1 all wait times172        # will be equal to `wait`. For any number of retries the last wait time will be173        # equal to `wait` and for retries>2 the first wait time will be equal to `wait / factor`.174        wait_times = iter(175            factor ** (n / n_waits - 1) * wait_time for n in range(0, n_waits + 1)176        )177 178        for attempt in range(self.register_session_retries + 1):179            try:180                smbclient.register_session(181                    self.host,182                    username=self.username,183                    password=self.password,184                    port=self._port,185                    encrypt=self.encrypt,186                    connection_timeout=self.timeout,187                )188                return189            except (190                smbprotocol.exceptions.SMBAuthenticationError,191                smbprotocol.exceptions.LogonFailure,192            ):193                # These exceptions should not be repeated, as they clearly indicate194                # that the credentials are invalid and not a network issue.195                raise196            except ValueError as exc:197                if re.findall(r"\[Errno -\d+]", str(exc)):198                    # This exception is raised by the smbprotocol.transport:Tcp.connect199                    # and originates from socket.gaierror (OSError). These exceptions might200                    # be raised due to network instability. We will retry to connect.201                    retried_errors.append(exc)202                else:203                    # All another ValueError exceptions should be raised, as they are not204                    # related to network issues.205                    raise206            except Exception as exc:207                # Save the exception and retry to connect. This except might be dropped208                # in the future, once all exceptions suited for retry are identified.209                retried_errors.append(exc)210 211            if attempt < self.register_session_retries:212                time.sleep(next(wait_times))213 214        # Raise last exception to inform user about the connection issues.215        # Note: Should we use ExceptionGroup to raise all exceptions?216        raise retried_errors[-1]217 218    @classmethod219    def _strip_protocol(cls, path):220        return infer_storage_options(path)["path"]221 222    @staticmethod223    def _get_kwargs_from_urls(path):224        # smb://workgroup;user:password@host:port/share/folder/file.csv225        out = infer_storage_options(path)226        out.pop("path", None)227        out.pop("protocol", None)228        return out229 230    def mkdir(self, path, create_parents=True, **kwargs):231        wpath = _as_unc_path(self.host, path)232        if create_parents:233            smbclient.makedirs(wpath, exist_ok=False, port=self._port, **kwargs)234        else:235            smbclient.mkdir(wpath, port=self._port, **kwargs)236 237    def makedirs(self, path, exist_ok=False):238        if _share_has_path(path):239            wpath = _as_unc_path(self.host, path)240            smbclient.makedirs(wpath, exist_ok=exist_ok, port=self._port)241 242    def rmdir(self, path):243        if _share_has_path(path):244            wpath = _as_unc_path(self.host, path)245            smbclient.rmdir(wpath, port=self._port)246 247    def info(self, path, **kwargs):248        wpath = _as_unc_path(self.host, path)249        stats = smbclient.stat(wpath, port=self._port, **kwargs)250        if S_ISDIR(stats.st_mode):251            stype = "directory"252        elif S_ISLNK(stats.st_mode):253            stype = "link"254        else:255            stype = "file"256        res = {257            "name": path + "/" if stype == "directory" else path,258            "size": stats.st_size,259            "type": stype,260            "uid": stats.st_uid,261            "gid": stats.st_gid,262            "time": stats.st_atime,263            "mtime": stats.st_mtime,264        }265        return res266 267    def created(self, path):268        """Return the created timestamp of a file as a datetime.datetime"""269        wpath = _as_unc_path(self.host, path)270        stats = smbclient.stat(wpath, port=self._port)271        return datetime.datetime.fromtimestamp(stats.st_ctime, tz=datetime.timezone.utc)272 273    def modified(self, path):274        """Return the modified timestamp of a file as a datetime.datetime"""275        wpath = _as_unc_path(self.host, path)276        stats = smbclient.stat(wpath, port=self._port)277        return datetime.datetime.fromtimestamp(stats.st_mtime, tz=datetime.timezone.utc)278 279    def ls(self, path, detail=True, **kwargs):280        unc = _as_unc_path(self.host, path)281        listed = smbclient.listdir(unc, port=self._port, **kwargs)282        dirs = ["/".join([path.rstrip("/"), p]) for p in listed]283        if detail:284            dirs = [self.info(d) for d in dirs]285        return dirs286 287    # pylint: disable=too-many-arguments288    def _open(289        self,290        path,291        mode="rb",292        block_size=-1,293        autocommit=True,294        cache_options=None,295        **kwargs,296    ):297        """298        block_size: int or None299            If 0, no buffering, 1, line buffering, >1, buffer that many bytes300 301        Notes302        -----303        By specifying 'share_access' in 'kwargs' it is possible to override the304        default shared access setting applied in the constructor of this object.305        """306        if self.auto_mkdir and "w" in mode:307            self.makedirs(self._parent(path), exist_ok=True)308        bls = block_size if block_size is not None and block_size >= 0 else -1309        wpath = _as_unc_path(self.host, path)310        share_access = kwargs.pop("share_access", self.share_access)311        if "w" in mode and autocommit is False:312            temp = _as_temp_path(self.host, path, self.temppath)313            return SMBFileOpener(314                wpath, temp, mode, port=self._port, block_size=bls, **kwargs315            )316        return smbclient.open_file(317            wpath,318            mode,319            buffering=bls,320            share_access=share_access,321            port=self._port,322            **kwargs,323        )324 325    def copy(self, path1, path2, **kwargs):326        """Copy within two locations in the same filesystem"""327        wpath1 = _as_unc_path(self.host, path1)328        wpath2 = _as_unc_path(self.host, path2)329        if self.auto_mkdir:330            self.makedirs(self._parent(path2), exist_ok=True)331        smbclient.copyfile(wpath1, wpath2, port=self._port, **kwargs)332 333    def _rm(self, path):334        if _share_has_path(path):335            wpath = _as_unc_path(self.host, path)336            stats = smbclient.stat(wpath, port=self._port)337            if S_ISDIR(stats.st_mode):338                smbclient.rmdir(wpath, port=self._port)339            else:340                smbclient.remove(wpath, port=self._port)341 342    def mv(self, path1, path2, recursive=None, maxdepth=None, **kwargs):343        wpath1 = _as_unc_path(self.host, path1)344        wpath2 = _as_unc_path(self.host, path2)345        smbclient.rename(wpath1, wpath2, port=self._port, **kwargs)346 347 348def _as_unc_path(host, path):349    rpath = path.replace("/", "\\")350    unc = f"\\\\{host}{rpath}"351    return unc352 353 354def _as_temp_path(host, path, temppath):355    share = path.split("/")[1]356    temp_file = f"/{share}{temppath}/{uuid.uuid4()}"357    unc = _as_unc_path(host, temp_file)358    return unc359 360 361def _share_has_path(path):362    parts = path.count("/")363    if path.endswith("/"):364        return parts > 2365    return parts > 1366 367 368class SMBFileOpener:369    """writes to remote temporary file, move on commit"""370 371    def __init__(self, path, temp, mode, port=445, block_size=-1, **kwargs):372        self.path = path373        self.temp = temp374        self.mode = mode375        self.block_size = block_size376        self.kwargs = kwargs377        self.smbfile = None378        self._incontext = False379        self.port = port380        self._open()381 382    def _open(self):383        if self.smbfile is None or self.smbfile.closed:384            self.smbfile = smbclient.open_file(385                self.temp,386                self.mode,387                port=self.port,388                buffering=self.block_size,389                **self.kwargs,390            )391 392    def commit(self):393        """Move temp file to definitive on success."""394        # TODO: use transaction support in SMB protocol395        smbclient.replace(self.temp, self.path, port=self.port)396 397    def discard(self):398        """Remove the temp file on failure."""399        smbclient.remove(self.temp, port=self.port)400 401    def __fspath__(self):402        return self.path403 404    def __iter__(self):405        return self.smbfile.__iter__()406 407    def __getattr__(self, item):408        return getattr(self.smbfile, item)409 410    def __enter__(self):411        self._incontext = True412        return self.smbfile.__enter__()413 414    def __exit__(self, exc_type, exc_value, traceback):415        self._incontext = False416        self.smbfile.__exit__(exc_type, exc_value, traceback)417