codekingpro/portable-devtools

1import os2import ssl3import uuid4from ftplib import FTP, FTP_TLS, Error, error_perm5from typing import Any6 7from ..spec import AbstractBufferedFile, AbstractFileSystem8from ..utils import infer_storage_options, isfilelike9 10SECURITY_PROTOCOL_MAP = {11    "tls": ssl.PROTOCOL_TLS,12    "tlsv1": ssl.PROTOCOL_TLSv1,13    "tlsv1_1": ssl.PROTOCOL_TLSv1_1,14    "tlsv1_2": ssl.PROTOCOL_TLSv1_2,15    "sslv23": ssl.PROTOCOL_SSLv23,16}17 18 19class ImplicitFTPTLS(FTP_TLS):20    """21    FTP_TLS subclass that automatically wraps sockets in SSL22    to support implicit FTPS.23    """24 25    def __init__(self, *args, **kwargs):26        super().__init__(*args, **kwargs)27        self._sock = None28 29    @property30    def sock(self):31        """Return the socket."""32        return self._sock33 34    @sock.setter35    def sock(self, value):36        """When modifying the socket, ensure that it is ssl wrapped."""37        if value is not None and not isinstance(value, ssl.SSLSocket):38            value = self.context.wrap_socket(value)39        self._sock = value40 41 42class FTPFileSystem(AbstractFileSystem):43    """A filesystem over classic FTP"""44 45    root_marker = "/"46    cachable = False47    protocol = "ftp"48 49    def __init__(50        self,51        host,52        port=21,53        username=None,54        password=None,55        acct=None,56        block_size=None,57        tempdir=None,58        timeout=30,59        encoding="utf-8",60        tls=False,61        **kwargs,62    ):63        """64        You can use _get_kwargs_from_urls to get some kwargs from65        a reasonable FTP url.66 67        Authentication will be anonymous if username/password are not68        given.69 70        Parameters71        ----------72        host: str73            The remote server name/ip to connect to74        port: int75            Port to connect with76        username: str or None77            If authenticating, the user's identifier78        password: str of None79            User's password on the server, if using80        acct: str or None81            Some servers also need an "account" string for auth82        block_size: int or None83            If given, the read-ahead or write buffer size.84        tempdir: str85            Directory on remote to put temporary files when in a transaction86        timeout: int87            Timeout of the ftp connection in seconds88        encoding: str89            Encoding to use for directories and filenames in FTP connection90        tls: bool or str91            Enable FTP-TLS for secure connections:92                - False: Plain FTP (default)93                - True: Explicit TLS (FTPS with AUTH TLS command)94                - "tls": Auto-negotiate highest protocol95                - "tlsv1": TLS v1.096                - "tlsv1_1": TLS v1.197                - "tlsv1_2": TLS v1.298        """99        super().__init__(**kwargs)100        self.host = host101        self.port = port102        self.tempdir = tempdir or "/tmp"103        self.cred = username or "", password or "", acct or ""104        self.timeout = timeout105        self.encoding = encoding106        if block_size is not None:107            self.blocksize = block_size108        else:109            self.blocksize = 2**16110        self.tls = tls111        self._connect()112        if isinstance(self.tls, bool) and self.tls:113            self.ftp.prot_p()114 115    def _connect(self):116        security = None117        if self.tls:118            if isinstance(self.tls, str):119                ftp_cls = ImplicitFTPTLS120                security = SECURITY_PROTOCOL_MAP.get(121                    self.tls,122                    f"Not supported {self.tls} protocol",123                )124                if isinstance(security, str):125                    raise ValueError(security)126            else:127                ftp_cls = FTP_TLS128        else:129            ftp_cls = FTP130        self.ftp = ftp_cls(timeout=self.timeout, encoding=self.encoding)131        if security:132            self.ftp.ssl_version = security133        self.ftp.connect(self.host, self.port)134        self.ftp.login(*self.cred)135 136    @classmethod137    def _strip_protocol(cls, path):138        return "/" + infer_storage_options(path)["path"].lstrip("/").rstrip("/")139 140    @staticmethod141    def _get_kwargs_from_urls(urlpath):142        out = infer_storage_options(urlpath)143        out.pop("path", None)144        out.pop("protocol", None)145        return out146 147    def ls(self, path, detail=True, **kwargs):148        path = self._strip_protocol(path)149        out = []150        if path not in self.dircache:151            try:152                try:153                    out = [154                        (fn, details)155                        for (fn, details) in self.ftp.mlsd(path)156                        if fn not in [".", ".."]157                        and details["type"] not in ["pdir", "cdir"]158                    ]159                except error_perm:160                    out = _mlsd2(self.ftp, path)  # Not platform independent161                for fn, details in out:162                    details["name"] = "/".join(163                        ["" if path == "/" else path, fn.lstrip("/")]164                    )165                    if details["type"] == "file":166                        details["size"] = int(details["size"])167                    else:168                        details["size"] = 0169                    if details["type"] == "dir":170                        details["type"] = "directory"171                self.dircache[path] = out172            except Error:173                try:174                    info = self.info(path)175                    if info["type"] == "file":176                        out = [(path, info)]177                except (Error, IndexError) as exc:178                    raise FileNotFoundError(path) from exc179        files = self.dircache.get(path, out)180        if not detail:181            return sorted([fn for fn, details in files])182        return [details for fn, details in files]183 184    def info(self, path, **kwargs):185        # implement with direct method186        path = self._strip_protocol(path)187        if path == "/":188            # special case, since this dir has no real entry189            return {"name": "/", "size": 0, "type": "directory"}190        files = self.ls(self._parent(path).lstrip("/"), True)191        try:192            out = next(f for f in files if f["name"] == path)193        except StopIteration as exc:194            raise FileNotFoundError(path) from exc195        return out196 197    def get_file(self, rpath, lpath, **kwargs):198        if self.isdir(rpath):199            if not os.path.exists(lpath):200                os.mkdir(lpath)201            return202        if isfilelike(lpath):203            outfile = lpath204        else:205            outfile = open(lpath, "wb")206 207        def cb(x):208            outfile.write(x)209 210        self.ftp.retrbinary(211            f"RETR {rpath}",212            blocksize=self.blocksize,213            callback=cb,214        )215        if not isfilelike(lpath):216            outfile.close()217 218    def cat_file(self, path, start=None, end=None, **kwargs):219        if end is not None:220            return super().cat_file(path, start, end, **kwargs)221        out = []222 223        def cb(x):224            out.append(x)225 226        try:227            self.ftp.retrbinary(228                f"RETR {path}",229                blocksize=self.blocksize,230                rest=start,231                callback=cb,232            )233        except (Error, error_perm) as orig_exc:234            raise FileNotFoundError(path) from orig_exc235        return b"".join(out)236 237    def _open(238        self,239        path,240        mode="rb",241        block_size=None,242        cache_options=None,243        autocommit=True,244        **kwargs,245    ):246        path = self._strip_protocol(path)247        block_size = block_size or self.blocksize248        return FTPFile(249            self,250            path,251            mode=mode,252            block_size=block_size,253            tempdir=self.tempdir,254            autocommit=autocommit,255            cache_options=cache_options,256        )257 258    def _rm(self, path):259        path = self._strip_protocol(path)260        self.ftp.delete(path)261        self.invalidate_cache(self._parent(path))262 263    def rm(self, path, recursive=False, maxdepth=None):264        paths = self.expand_path(path, recursive=recursive, maxdepth=maxdepth)265        for p in reversed(paths):266            if self.isfile(p):267                self.rm_file(p)268            else:269                self.rmdir(p)270 271    def mkdir(self, path: str, create_parents: bool = True, **kwargs: Any) -> None:272        path = self._strip_protocol(path)273        parent = self._parent(path)274        if parent != self.root_marker and not self.exists(parent) and create_parents:275            self.mkdir(parent, create_parents=create_parents)276 277        self.ftp.mkd(path)278        self.invalidate_cache(self._parent(path))279 280    def makedirs(self, path: str, exist_ok: bool = False) -> None:281        path = self._strip_protocol(path)282        if self.exists(path):283            # NB: "/" does not "exist" as it has no directory entry284            if not exist_ok:285                raise FileExistsError(f"{path} exists without `exist_ok`")286            # exists_ok=True -> no-op287        else:288            self.mkdir(path, create_parents=True)289 290    def rmdir(self, path):291        path = self._strip_protocol(path)292        self.ftp.rmd(path)293        self.invalidate_cache(self._parent(path))294 295    def mv(self, path1, path2, **kwargs):296        path1 = self._strip_protocol(path1)297        path2 = self._strip_protocol(path2)298        self.ftp.rename(path1, path2)299        self.invalidate_cache(self._parent(path1))300        self.invalidate_cache(self._parent(path2))301 302    def __del__(self):303        self.ftp.close()304 305    def invalidate_cache(self, path=None):306        if path is None:307            self.dircache.clear()308        else:309            self.dircache.pop(path, None)310        super().invalidate_cache(path)311 312 313class TransferDone(Exception):314    """Internal exception to break out of transfer"""315 316    pass317 318 319class FTPFile(AbstractBufferedFile):320    """Interact with a remote FTP file with read/write buffering"""321 322    def __init__(323        self,324        fs,325        path,326        mode="rb",327        block_size="default",328        autocommit=True,329        cache_type="readahead",330        cache_options=None,331        **kwargs,332    ):333        super().__init__(334            fs,335            path,336            mode=mode,337            block_size=block_size,338            autocommit=autocommit,339            cache_type=cache_type,340            cache_options=cache_options,341            **kwargs,342        )343        if not autocommit:344            self.target = self.path345            self.path = "/".join([kwargs["tempdir"], str(uuid.uuid4())])346 347    def commit(self):348        self.fs.mv(self.path, self.target)349 350    def discard(self):351        self.fs.rm(self.path)352 353    def _fetch_range(self, start, end):354        """Get bytes between given byte limits355 356        Implemented by raising an exception in the fetch callback when the357        number of bytes received reaches the requested amount.358 359        Will fail if the server does not respect the REST command on360        retrieve requests.361        """362        out = []363        total = [0]364 365        def callback(x):366            total[0] += len(x)367            if total[0] > end - start:368                out.append(x[: (end - start) - total[0]])369                if end < self.size:370                    raise TransferDone371            else:372                out.append(x)373 374            if total[0] == end - start and end < self.size:375                raise TransferDone376 377        try:378            self.fs.ftp.retrbinary(379                f"RETR {self.path}",380                blocksize=self.blocksize,381                rest=start,382                callback=callback,383            )384        except TransferDone:385            try:386                # stop transfer, we got enough bytes for this block387                self.fs.ftp.abort()388                self.fs.ftp.getmultiline()389            except Error:390                self.fs._connect()391 392        return b"".join(out)393 394    def _upload_chunk(self, final=False):395        self.buffer.seek(0)396        self.fs.ftp.storbinary(397            f"STOR {self.path}", self.buffer, blocksize=self.blocksize, rest=self.offset398        )399        return True400 401 402def _mlsd2(ftp, path="."):403    """404    Fall back to using `dir` instead of `mlsd` if not supported.405 406    This parses a Linux style `ls -l` response to `dir`, but the response may407    be platform dependent.408 409    Parameters410    ----------411    ftp: ftplib.FTP412    path: str413        Expects to be given path, but defaults to ".".414    """415    lines = []416    minfo = []417    ftp.dir(path, lines.append)418    for line in lines:419        split_line = line.split()420        if len(split_line) < 9:421            continue422        this = (423            split_line[-1],424            {425                "modify": " ".join(split_line[5:8]),426                "unix.owner": split_line[2],427                "unix.group": split_line[3],428                "unix.mode": split_line[0],429                "size": split_line[4],430            },431        )432        if this[1]["unix.mode"][0] == "d":433            this[1]["type"] = "dir"434        else:435            this[1]["type"] = "file"436        minfo.append(this)437    return minfo438