lakehouse_engine.utils.extraction.sftp_extraction_utils

  1"""Utilities module for SFTP extraction processes."""
  2
  3import stat
  4from base64 import decodebytes
  5from datetime import datetime
  6from enum import Enum
  7from logging import Logger
  8from stat import S_ISREG
  9from typing import Any, List, Set, Tuple, Union
 10
 11import paramiko as p
 12from paramiko import Ed25519Key, PKey, RSAKey, Transport
 13from paramiko.sftp_client import SFTPAttributes, SFTPClient  # type: ignore
 14
 15from lakehouse_engine.transformers.exceptions import WrongArgumentsException
 16from lakehouse_engine.utils.logging_handler import LoggingHandler
 17
 18
 19class SFTPInputFormat(Enum):
 20    """Formats of algorithm input."""
 21
 22    CSV = "csv"
 23    FWF = "fwf"
 24    JSON = "json"
 25    XML = "xml"
 26
 27
 28class SFTPExtractionFilter(Enum):
 29    """Standardize the types of filters we can have from a SFTP source."""
 30
 31    file_name_contains = "file_name_contains"
 32    LATEST_FILE = "latest_file"
 33    EARLIEST_FILE = "earliest_file"
 34    GREATER_THAN = "date_time_gt"
 35    LOWER_THAN = "date_time_lt"
 36
 37
 38class SFTPExtractionUtils(object):
 39    """Utils for managing data extraction from particularly relevant SFTP sources."""
 40
 41    _logger: Logger = LoggingHandler(__name__).get_logger()
 42
 43    @classmethod
 44    def get_files_list(
 45        cls, sftp: SFTPClient, remote_path: str, options_args: dict
 46    ) -> Set[str]:
 47        """Get a list of files to be extracted from SFTP.
 48
 49        The arguments (options_args) to list files are:
 50        - date_time_gt(str):
 51            Filter the files greater than the string datetime
 52            formatted as "YYYY-MM-DD" or "YYYY-MM-DD HH:MM:SS".
 53        - date_time_lt(str):
 54            Filter the files lower than the string datetime
 55            formatted as "YYYY-MM-DD" or "YYYY-MM-DD HH:MM:SS".
 56        - earliest_file(bool):
 57            Filter the earliest dated file in the directory.
 58        - file_name_contains(str):
 59            Filter files when match the pattern.
 60        - latest_file(bool):
 61            Filter the most recent dated file in the directory.
 62        - sub_dir(bool):
 63            When true, the engine will search files into subdirectories
 64            of the remote_path.
 65            It will consider one level below the remote_path.
 66            When sub_dir is used with latest_file/earliest_file argument,
 67            the engine will retrieve the latest_file/earliest_file
 68            for each subdirectory.
 69
 70        Args:
 71            sftp: the SFTP client object.
 72            remote_path: path of files to be filtered.
 73            options_args: options from the acon.
 74
 75        Returns:
 76            A list containing the file names to be passed to Spark.
 77        """
 78        all_items, folder_path = cls._get_folder_items(remote_path, sftp, options_args)
 79
 80        filtered_files: Set[str] = set()
 81
 82        try:
 83            for item, folder in zip(all_items, folder_path):
 84                file_contains = cls._file_has_pattern(item, options_args)
 85                file_in_interval = cls._file_in_date_interval(item, options_args)
 86                if file_contains and file_in_interval:
 87                    filtered_files.add(folder + item.filename)
 88
 89            if (
 90                SFTPExtractionFilter.EARLIEST_FILE.value in options_args.keys()
 91                or SFTPExtractionFilter.LATEST_FILE.value in options_args.keys()
 92            ):
 93                filtered_files = cls._get_earliest_latest_file(
 94                    sftp, options_args, filtered_files, folder_path
 95                )
 96
 97        except Exception as e:
 98            cls._logger.error(f"SFTP list_files EXCEPTION: - {e}")
 99        return filtered_files
100
101    @classmethod
102    def get_sftp_client(
103        cls,
104        options_args: dict,
105    ) -> Tuple[SFTPClient, Transport]:
106        """Get the SFTP client.
107
108        The SFTP client is used to open an SFTP session across an open
109        SSH Transport and perform remote file operations.
110
111        Args:
112            options_args: dictionary containing SFTP connection parameters.
113                The Paramiko arguments expected to connect are:
114                    - "hostname": the server to connect to.
115                    - "port": the server port to connect to.
116                    - "username": the username to authenticate as.
117                    - "password": used for password authentication.
118                    - "pkey": optional - an optional public key to use for
119                        authentication.
120                    - "passphrase" – optional - options used for decrypting private
121                        keys.
122                    - "key_filename" – optional - the filename, or list of filenames,
123                        of optional private key(s) and/or certs to try for
124                        authentication.
125                    - "timeout" – an optional timeout (in seconds) for the TCP connect.
126                    - "allow_agent" – optional - set to False to disable
127                        connecting to the SSH agent.
128                    - "look_for_keys" – optional - set to False to disable searching
129                        for discoverable private key files in ~/.ssh/.
130                    - "compress" – optional - set to True to turn on compression.
131                    - "sock" - optional - an open socket or socket-like object
132                        to use for communication to the target host.
133                    - "gss_auth" – optional - True if you want to use GSS-API
134                        authentication.
135                    - "gss_kex" – optional - Perform GSS-API Key Exchange and
136                        user authentication.
137                    - "gss_deleg_creds" – optional - Delegate GSS-API client
138                        credentials or not.
139                    - "gss_host" – optional - The targets name in the kerberos database.
140                    - "gss_trust_dns" – optional - Indicates whether or
141                        not the DNS is trusted to securely canonicalize the name of the
142                        host being connected to (default True).
143                    - "banner_timeout" – an optional timeout (in seconds)
144                        to wait for the SSH banner to be presented.
145                    - "auth_timeout" – an optional timeout (in seconds)
146                        to wait for an authentication response.
147                    - "disabled_algorithms" – an optional dict passed directly to
148                        Transport and its keyword argument of the same name.
149                    - "transport_factory" – an optional callable which is handed a
150                        subset of the constructor arguments (primarily those related
151                        to the socket, GSS functionality, and algorithm selection)
152                        and generates a Transport instance to be used by this client.
153                        Defaults to Transport.__init__.
154
155                The parameter to specify the private key is expected to be in
156                RSA format. Attempting a connection with a blank host key is
157                not allowed unless the argument "add_auto_policy" is explicitly
158                set to True.
159
160        Returns:
161            sftp -> a new SFTPClient session object.
162            transport -> the Transport for this connection.
163        """
164        ssh_client = p.SSHClient()
165        try:
166            if not options_args.get("pkey") and not options_args.get("add_auto_policy"):
167                raise WrongArgumentsException(
168                    "Get SFTP Client: No host key (pkey) was provided and the "
169                    + "add_auto_policy property is false."
170                )
171
172            if options_args.get("pkey") and not options_args.get("key_type"):
173                raise WrongArgumentsException(
174                    "Get SFTP Client: The key_type must be provided when "
175                    + "the host key (pkey) is provided."
176                )
177
178            if options_args.get("pkey", None) and options_args.get("key_type", None):
179                key = cls._get_host_keys(
180                    options_args.get("pkey", None), options_args.get("key_type", None)
181                )
182                ssh_client.get_host_keys().add(
183                    hostname=f"[{options_args.get('hostname')}]:"
184                    + f"{options_args.get('port')}",
185                    keytype="ssh-rsa",
186                    key=key,
187                )
188            elif options_args.get("add_auto_policy", None):
189                ssh_client.load_system_host_keys()
190                ssh_client.set_missing_host_key_policy(p.WarningPolicy())  # nosec: B507
191            else:
192                ssh_client.load_system_host_keys()
193                ssh_client.set_missing_host_key_policy(p.RejectPolicy())
194
195            ssh_client.connect(
196                hostname=options_args.get("hostname"),
197                port=options_args.get("port", 22),
198                username=options_args.get("username", None),
199                password=options_args.get("password", None),
200                key_filename=options_args.get("key_filename", None),
201                timeout=options_args.get("timeout", None),
202                allow_agent=options_args.get("allow_agent", True),
203                look_for_keys=options_args.get("look_for_keys", True),
204                compress=options_args.get("compress", False),
205                sock=options_args.get("sock", None),
206                gss_auth=options_args.get("gss_auth", False),
207                gss_kex=options_args.get("gss_kex", False),
208                gss_deleg_creds=options_args.get("gss_deleg_creds", False),
209                gss_host=options_args.get("gss_host", False),
210                banner_timeout=options_args.get("banner_timeout", None),
211                auth_timeout=options_args.get("auth_timeout", None),
212                gss_trust_dns=options_args.get("gss_trust_dns", None),
213                passphrase=options_args.get("passphrase", None),
214                disabled_algorithms=options_args.get("disabled_algorithms", None),
215                transport_factory=options_args.get("transport_factory", None),
216            )
217
218            sftp = ssh_client.open_sftp()
219            transport = ssh_client.get_transport()
220        except ConnectionError as e:
221            cls._logger.error(e)
222            raise
223        return sftp, transport
224
225    @classmethod
226    def validate_format(cls, files_format: str) -> str:
227        """Validate the file extension based on the format definitions.
228
229        Args:
230            files_format: a string containing the file extension.
231
232        Returns:
233            The string validated and formatted.
234        """
235        formats_allowed = [
236            SFTPInputFormat.CSV.value,
237            SFTPInputFormat.FWF.value,
238            SFTPInputFormat.JSON.value,
239            SFTPInputFormat.XML.value,
240        ]
241
242        if files_format not in formats_allowed:
243            raise WrongArgumentsException(
244                f"The formats allowed for SFTP are {formats_allowed}."
245            )
246
247        return files_format
248
249    @classmethod
250    def validate_location(cls, location: str) -> str:
251        """Validate the location. Add "/" in the case it does not exist.
252
253        Args:
254            location: file path.
255
256        Returns:
257            The location validated.
258        """
259        return location if location.rfind("/") == len(location) - 1 else location + "/"
260
261    @classmethod
262    def _file_has_pattern(cls, item: SFTPAttributes, options_args: dict) -> bool:
263        """Check if a file follows the pattern used for filtering.
264
265        Args:
266            item: item available in SFTP directory.
267            options_args: options from the acon.
268
269        Returns:
270            A boolean telling whether the file contains a pattern or not.
271        """
272        file_to_consider = True
273
274        if SFTPExtractionFilter.file_name_contains.value in options_args.keys():
275            if not (
276                options_args.get(SFTPExtractionFilter.file_name_contains.value)
277                in item.filename
278                and (S_ISREG(item.st_mode) or cls._is_compressed(item.filename))
279            ):
280                file_to_consider = False
281
282        return file_to_consider
283
284    @classmethod
285    def _file_in_date_interval(
286        cls,
287        item: SFTPAttributes,
288        options_args: dict,
289    ) -> bool:
290        """Check if the file is in the expected date interval.
291
292        The logic is applied based on the arguments greater_than and lower_than.
293        i.e:
294        - if greater_than and lower_than have values,
295        then it performs a between.
296        - if only lower_than has values,
297        then only values lower than the input value will be retrieved.
298        - if only greater_than has values,
299        then only values greater than the input value will be retrieved.
300
301        Args:
302            item: item available in SFTP directory.
303            options_args: options from the acon.
304
305        Returns:
306            A boolean telling whether the file is in the expected date interval or not.
307        """
308        file_to_consider = True
309
310        if (
311            SFTPExtractionFilter.LOWER_THAN.value in options_args.keys()
312            or SFTPExtractionFilter.GREATER_THAN.value in options_args.keys()
313            and (S_ISREG(item.st_mode) or cls._is_compressed(item.filename))
314        ):
315            lower_than = options_args.get(
316                SFTPExtractionFilter.LOWER_THAN.value, "9999-12-31"
317            )
318            greater_than = options_args.get(
319                SFTPExtractionFilter.GREATER_THAN.value, "1900-01-01"
320            )
321
322            file_date = datetime.fromtimestamp(item.st_mtime)
323
324            if not (
325                (
326                    lower_than == greater_than
327                    and cls._validate_date(greater_than)
328                    <= file_date
329                    <= cls._validate_date(lower_than)
330                )
331                or (
332                    cls._validate_date(greater_than)
333                    < file_date
334                    < cls._validate_date(lower_than)
335                )
336            ):
337                file_to_consider = False
338
339        return file_to_consider
340
341    @classmethod
342    def _get_earliest_latest_file(
343        cls,
344        sftp: SFTPClient,
345        options_args: dict,
346        list_filter_files: Set[str],
347        folder_path: List,
348    ) -> Set[str]:
349        """Get the earliest or latest file of a directory.
350
351        Args:
352            sftp: the SFTP client object.
353            options_args: options from the acon.
354            list_filter_files: set of file names to filter from.
355            folder_path: the location of files.
356
357        Returns:
358            A set containing the earliest/latest file name.
359        """
360        list_earl_lat_files: Set[str] = set()
361
362        for folder in folder_path:
363            file_date = 0
364            file_name = ""
365            all_items, file_path = cls._get_folder_items(
366                f"{folder}", sftp, options_args
367            )
368            for item in all_items:
369                if (
370                    folder + item.filename in list_filter_files
371                    and (S_ISREG(item.st_mode) or cls._is_compressed(item.filename))
372                    and (
373                        options_args.get("earliest_file")
374                        and (file_date == 0 or item.st_mtime < file_date)
375                    )
376                    or (
377                        options_args.get("latest_file")
378                        and (file_date == 0 or item.st_mtime > file_date)
379                    )
380                ):
381                    file_date = item.st_mtime
382                    file_name = folder + item.filename
383            list_earl_lat_files.add(file_name)
384
385        return list_earl_lat_files
386
387    @classmethod
388    def _get_folder_items(
389        cls, remote_path: str, sftp: SFTPClient, options_args: dict
390    ) -> Tuple:
391        """Get the files and the directory to be processed.
392
393        Args:
394            remote_path: root folder path.
395            sftp: a SFTPClient session object.
396            options_args: options from the acon.
397
398        Returns:
399            A tuple with a list of items (file object) and a list of directories.
400        """
401        sub_dir = options_args.get("sub_dir", False)
402        all_items: List[SFTPAttributes] = sftp.listdir_attr(remote_path)
403        items: List[SFTPAttributes] = []
404        folders: List = []
405
406        for item in all_items:
407            is_dir = stat.S_ISDIR(item.st_mode)
408            if is_dir and sub_dir and not item.filename.endswith((".gz", ".zip")):
409                dirs = sftp.listdir_attr(f"{remote_path}{item.filename}")
410                for file in dirs:
411                    items.append(file)
412                    folders.append(f"{remote_path}{item.filename}/")
413            else:
414                items.append(item)
415                folders.append(remote_path)
416
417        return items, folders
418
419    @classmethod
420    def _get_host_keys(cls, pkey: str, key_type: str) -> PKey:
421        """Get the pkey that will be added to the server.
422
423        Args:
424            pkey: a string with a host key value.
425            key_type: the type of key (rsa or ed25519).
426
427        Returns:
428            A PKey that will be used to authenticate the connection.
429        """
430        key: Union[RSAKey, Ed25519Key] = None
431        if pkey and key_type.lower() == "rsa":
432            b_pkey = bytes(pkey, "UTF-8")
433            key = p.RSAKey(data=decodebytes(b_pkey))
434        elif pkey and key_type.lower() == "ed25519":
435            b_pkey = bytes(pkey, "UTF-8")
436            key = p.Ed25519Key(data=decodebytes(b_pkey))
437
438        return key
439
440    @classmethod
441    def _is_compressed(cls, filename: str) -> Any:
442        """Validate if it is a compressed file.
443
444        Args:
445            filename: name of the file to be validated.
446
447        Returns:
448            A boolean with the result.
449        """
450        return filename.endswith((".gz", ".zip"))
451
452    @classmethod
453    def _validate_date(cls, date_text: str) -> datetime:
454        """Validate the input date format.
455
456        Args:
457            date_text: a string with the date or datetime value.
458            The expected formats are:
459                YYYY-MM-DD and YYYY-MM-DD HH:MM:SS
460
461        Returns:
462            The datetime validated and formatted.
463        """
464        for fmt in ("%Y-%m-%d", "%Y-%m-%d %H:%M:%S"):
465            try:
466                if date_text is not None:
467                    return datetime.strptime(date_text, fmt)
468            except ValueError:
469                pass
470        raise ValueError(
471            "Incorrect data format, should be YYYY-MM-DD or YYYY-MM-DD HH:MM:SS."
472        )