glasswall.utils

  1import ctypes as ct
  2import functools
  3import io
  4import math
  5import os
  6import pathlib
  7import stat
  8import tempfile
  9import warnings
 10from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 11
 12from lxml import etree
 13
 14import glasswall
 15from glasswall.config.logging import log
 16
 17
 18def as_bytes(file_: Union[bytes, bytearray, io.BytesIO]):
 19    """ Returns file_ as bytes.
 20
 21    Args:
 22        file_ (Union[bytes, bytearray, io.BytesIO]): The file
 23
 24    Returns:
 25        bytes
 26
 27    Raises:
 28        TypeError: If file_ is not an instance of: bytes, bytearray, io.BytesIO
 29    """
 30    if isinstance(file_, bytes):
 31        return file_
 32    elif isinstance(file_, bytearray):
 33        return bytes(file_)
 34    elif isinstance(file_, io.BytesIO):
 35        file_.seek(0)
 36        return file_.read()
 37    else:
 38        raise TypeError(file_)
 39
 40
 41def as_io_BytesIO(file_: Union[bytes, bytearray]):
 42    """ Returns file_ as io.BytesIO object.
 43
 44    Args:
 45        file_ (Union[bytes, bytearray]): The bytes or bytearray of the file
 46
 47    Returns:
 48        io.BytesIO object
 49
 50    Raises:
 51        TypeError: If file_ is not an instance of: bytes, bytearray, io.BytesIO
 52    """
 53    if isinstance(file_, bytes):
 54        return io.BytesIO(file_)
 55    elif isinstance(file_, bytearray):
 56        return io.BytesIO(bytes(file_))
 57    elif isinstance(file_, io.BytesIO):
 58        return file_
 59    else:
 60        raise TypeError(file_)
 61
 62
 63def as_snake_case(string):
 64    return ''.join(
 65        [
 66            '_' + char.lower()
 67            if char.isupper() else char
 68            for char in string
 69        ]
 70    ).lstrip('_')
 71
 72
 73def as_title(string):
 74    return ''.join(
 75        word.title()
 76        for word in string.replace(".", "_").split("_")
 77    )
 78
 79
 80def buffer_to_bytes(buffer: ct.c_void_p, buffer_length: ct.c_size_t):
 81    """ Convert ctypes buffer and buffer_length to bytes.
 82
 83    Args:
 84        buffer (ct.c_void_p()): The file buffer.
 85        buffer_length (ct.c_size_t()): The file buffer length.
 86
 87    Returns:
 88        bytes (bytes): The file as bytes.
 89    """
 90
 91    file_buffer = (ct.c_byte * buffer_length.value)()
 92    ct.memmove(file_buffer, buffer.value, buffer_length.value)
 93
 94    return bytes(file_buffer)
 95
 96
 97class CwdHandler:
 98    """ Changes the current working directory to new_cwd on __enter__, and back to previous cwd on __exit__.
 99
100    Args:
101        new_cwd (str): The new current working directory to temporarily change to.
102    """
103
104    def __init__(self, new_cwd: str):
105        self.new_cwd = new_cwd if os.path.isdir(new_cwd) else os.path.dirname(new_cwd)
106        self.old_cwd = os.getcwd()
107
108    def __enter__(self):
109        os.chdir(self.new_cwd)
110
111    def __exit__(self, type, value, traceback):
112        os.chdir(self.old_cwd)
113
114
115def delete_directory(directory: str, keep_folder: bool = False):
116    """ Delete a directory and its contents.
117
118    Args:
119        directory (str): The directory path.
120        keep_folder (bool, optional): Default False. If False, only delete contents.
121    """
122    if os.path.isdir(directory):
123        # Delete all files in directory
124        for file_ in list_file_paths(directory, followlinks=False):
125            os.remove(file_)
126
127        # Delete all empty subdirectories
128        delete_empty_subdirectories(directory)
129
130        # Delete the directory
131        if keep_folder is False:
132            os.rmdir(directory)
133
134
135def delete_empty_subdirectories(directory: str):
136    """ Deletes all empty subdirectories of a given directory.
137
138    Args:
139        directory (str): The directory to delete subdirectories from.
140
141    Returns:
142        None
143    """
144
145    for root, dirs, _ in os.walk(directory, topdown=False):
146        for dir_ in dirs:
147            absolute_path = os.path.join(root, dir_)
148            try:
149                os.rmdir(absolute_path)
150            except PermissionError:
151                # directory might be read-only
152                try:
153                    os.chmod(absolute_path, stat.S_IWRITE)
154                except Exception:
155                    log.warning(f"PermissionError while attempting to delete {absolute_path}. Attempted chmod but failed.")
156                try:
157                    os.rmdir(absolute_path)
158                except OSError:
159                    # cannot be deleted
160                    pass
161            except OSError:
162                # not empty, don't delete
163                pass
164
165
166def flatten_list(list_: Iterable):
167    """ Returns a flattened list. [[1, 2], ["3"], (4, 5,), [6]] --> [1, 2, "3", 4, 5, 6] """
168    return [
169        item
170        for sublist in list_
171        for item in sublist
172    ]
173
174
175def get_file_type(file_path: str):
176    """ Returns the filetype of a file. "data/files/splat.zip" -> "zip" """
177    return os.path.splitext(file_path)[-1].replace(".", "")
178
179
180def get_libraries(directory: str, library_names: Optional[List[str]] = None, ignore_errors: bool = False):
181    """ Recursively calls get_library on each library from glasswall.libraries.os_info on the given directory.
182
183    Args:
184        directory (str): The directory to search from.
185        library_names (List[str], optional): List of libraries to return, if None iterates all libraries found in glasswall.libraries.os_info
186        ignore_errors (bool, optional): Default False, prevents get_library raising FileNotFoundError when True.
187
188    Returns:
189        libraries (dict[str, str]): A dictionary of library names and their absolute file paths.
190    """
191    libraries = {}
192
193    if not library_names:
194        library_names = glasswall.libraries.os_info[glasswall._OPERATING_SYSTEM].keys()
195
196    for library_name in library_names:
197        try:
198            libraries[library_name] = get_library(library_name, directory)
199        except FileNotFoundError:
200            if ignore_errors is True:
201                continue
202            raise
203
204    return libraries
205
206
207def get_library(library: str, directory: str):
208    """ Returns a path to the specified library found from the current directory or any subdirectory. If multiple libraries exist, returns the file with the latest modified time.
209
210    Args:
211        library (str): The library to search for, ie: "rebuild", "word_search"
212        directory (str): The directory to search from.
213
214    Returns:
215        library_file_path (str): The absolute file path to the library.
216
217    Raises:
218        KeyError: Unsupported OS or library name was not found in glasswall.libraries.os_info.
219        FileNotFoundError: Library was not found.
220    """
221    if not os.path.isdir(directory):
222        raise NotADirectoryError(directory)
223
224    library = as_snake_case(library)
225    library_file_names = glasswall.libraries.os_info[glasswall._OPERATING_SYSTEM][library]["file_name"]
226
227    if isinstance(library_file_names, str):
228        library_file_names = [library_file_names]
229
230    matches = []
231    for alias in library_file_names:
232        p = pathlib.Path(directory)
233        alias_matches = list(p.rglob(alias))
234        matches.extend(alias_matches)
235
236    if matches:
237        latest_library = str(max(matches, key=os.path.getctime).resolve())
238        if len(matches) > 1:
239            # warn that multiple libraries found, list library paths if there are <= 5
240            if len(matches) <= 5:
241                log.warning(f"Found {len(matches)} {library} libraries, but expected only one:\n{chr(10).join(str(item) for item in matches)}\nLatest library: {latest_library}")
242            else:
243                log.warning(f"Found {len(matches)} {library} libraries, but expected only one.\nLatest library: {latest_library}")
244
245        # Return library with latest change time
246        return latest_library
247
248    # exhausted, not found
249    raise FileNotFoundError(f'Could not find any files: "{library_file_names}" under directory: "{directory}"')
250
251
252def iterate_directory_entries(directory: str, file_type: str = 'all', absolute: bool = True, recursive: bool = True, followlinks: bool = True, start_directory: str = None):
253    """ Generate entries (files, directories, or both) in a given directory using os.scandir().
254
255    Args:
256        directory (str): The path to the directory whose entries are to be listed.
257        file_type (str, optional): Type of entries to return.
258            - 'all': Return both files and directories (default).
259            - 'files': Return only files.
260            - 'directories': Return only directories.
261        absolute (bool, optional): Whether to return absolute paths (default) or relative paths.
262        recursive (bool, optional): Whether to recurse into subdirectories (default is True).
263        followlinks (bool, optional): Whether to follow symbolic links and yield entries from the target directory (default is True).
264        start_directory (str, optional): The starting directory used to calculate relative paths (default is None).
265
266    Yields:
267        str: The full path of each file or directory found in the specified directory.
268
269    Raises:
270        ValueError: If an invalid 'file_type' value is provided.
271        NotADirectoryError: If the directory does not exist.
272
273    Example:
274        directory = '/path/to/your/directory'
275
276        # Iterate through all entries (files and directories) in the directory
277        for entry in iterate_directory_entries(directory):
278            print(entry)
279
280        # Iterate through only file entries in the directory
281        for file in iterate_directory_entries(directory, file_type='files'):
282            print("File:", file)
283
284        # Iterate through only directory entries in the directory
285        for directory in iterate_directory_entries(directory, file_type='directories'):
286            print("Directory:", directory)
287    """
288    if not os.path.isdir(directory):
289        raise NotADirectoryError(directory)
290
291    allowed_types = ['all', 'files', 'directories']
292
293    # Check if the provided file_type is valid
294    if file_type not in allowed_types:
295        raise ValueError(f"Invalid file_type '{file_type}'. Allowed values are {', '.join(allowed_types)}.")
296
297    # Convert the directory to an absolute path
298    directory = os.path.abspath(directory)
299
300    # Set the start_directory to the provided directory if not specified
301    start_directory = start_directory or directory
302
303    # Get the directory entries using os.scandir()
304    for entry in os.scandir(directory):
305        # If the entry is a directory
306        if entry.is_dir(follow_symlinks=followlinks):
307            # If recursive is True, traverse the subdirectory
308            if recursive:
309                yield from iterate_directory_entries(entry.path, file_type, absolute, recursive, followlinks, start_directory)
310
311            # If the file_type is not "files", yield the directory entry
312            if file_type != "files":
313                if absolute:
314                    yield entry.path
315                else:
316                    yield os.path.relpath(entry.path, start=start_directory)
317
318        # If the entry is a file
319        elif entry.is_file(follow_symlinks=followlinks):
320            # If the file_type is not "directories", yield the file entry
321            if file_type != "directories":
322                if absolute:
323                    yield entry.path
324                else:
325                    yield os.path.relpath(entry.path, start=start_directory)
326
327
328def list_file_paths(directory: str, file_type: str = 'files', absolute: bool = True, recursive: bool = True, followlinks: bool = True) -> list:
329    """ List all file paths in a given directory and its subdirectories.
330
331    Args:
332        directory (str): The path to the directory whose file paths are to be listed.
333        file_type (str, optional): Type of entries to return.
334            - 'all': Return both files and directories.
335            - 'files': Return only files (default).
336            - 'directories': Return only directories.
337        absolute (bool, optional): Whether to return absolute paths (default is True).
338        recursive (bool, optional): Whether to recurse into subdirectories (default is True).
339        followlinks (bool, optional): Whether to follow symbolic links and list file paths from the target directory (default is True).
340
341    Returns:
342        list: A list of file paths found in the specified directory and its subdirectories.
343
344    Example:
345        directory = '/path/to/your/directory'
346        file_paths = list_file_paths(directory)
347        print(file_paths)
348    """
349    # Remove duplicate file paths (symlinks of same files or other symlinks), and sort
350    return sorted(set(iterate_directory_entries(directory, file_type, absolute, recursive, followlinks)))
351
352
353def list_subdirectory_paths(directory: str, recursive: bool = False, absolute: bool = True):
354    """ Returns a list of paths to subdirectories in a directory.
355
356    Args:
357        directory (str): The directory to list subdirectories from.
358        recursive (bool, optional): Default False. Include subdirectories of subdirectories.
359        absolute (bool, optional): Default True. Return paths as absolute paths. If False, returns relative paths.
360
361    Returns:
362        subdirectories (list): A list of subdirectory paths.
363    """
364    subdirectories = [f.path for f in os.scandir(directory) if f.is_dir()]
365
366    if recursive:
367        for subdirectory in subdirectories:
368            subdirectories.extend(list_subdirectory_paths(subdirectory, recursive=True))
369
370    if absolute:
371        subdirectories = [os.path.abspath(path) for path in subdirectories]
372    else:
373        subdirectories = [os.path.relpath(path, directory) for path in subdirectories]
374
375    return subdirectories
376
377
378def load_dependencies(dependencies: list, ignore_errors: bool = False):
379    """ Calls ctypes.cdll.LoadLibrary on each file path in `dependencies`.
380
381    Args:
382        dependencies (list): A list of absolute file paths of library dependencies.
383        ignore_errors (bool, optional): Default False, avoid raising exceptions from ct.cdll.LoadLibrary if ignore_errors is True.
384
385    Returns:
386        missing_dependencies (list): A list of missing dependencies, or an empty list.
387    """
388    missing_dependencies = [dependency for dependency in dependencies if not os.path.isfile(dependency)]
389
390    for dependency in dependencies:
391        # Try to load dependencies that exist
392        if dependency not in missing_dependencies:
393            try:
394                ct.cdll.LoadLibrary(dependency)
395            except Exception:
396                if ignore_errors:
397                    pass
398                else:
399                    raise
400
401    return missing_dependencies
402
403
404def round_up(number: float, decimals=0) -> float:
405    """ Rounds a number up to a specified number of decimal places.
406
407    Args:
408        number (float): The number to be rounded.
409        decimals (int, optional): The number of decimal places to round to. Defaults to 0.
410
411    Returns:
412        float: The rounded number.
413
414    Examples:
415        >>> round_up(105, 0)
416        105.0
417        >>> round_up(0.015, 2)
418        0.02
419        >>> round_up(0.025, 2)
420        0.03
421        >>> round_up(0.00001, 2)
422        0.01
423    """
424    multiplier = 10 ** decimals
425    return math.ceil(number * multiplier) / multiplier
426
427
428class TempDirectoryPath:
429    """ Gives a path to a uniquely named temporary directory that does not currently exist on __enter__, deletes the directory if it exists on __exit__.
430
431    Args:
432        delete (bool, optional): Default True. Delete the temporary directory on __exit__
433    """
434
435    def __init__(self, delete: bool = True):
436        # Validate args
437        if not isinstance(delete, bool):
438            raise TypeError(delete)
439
440        self.temp_directory = None
441        self.delete = delete
442
443        while self.temp_directory is None or os.path.isdir(self.temp_directory):
444            self.temp_directory = os.path.join(glasswall._TEMPDIR, next(tempfile._get_candidate_names()))
445
446        # Normalize
447        self.temp_directory = str(pathlib.Path(self.temp_directory).resolve())
448
449        # Create temp directory
450        os.makedirs(self.temp_directory, exist_ok=True)
451
452    def __enter__(self):
453        return self.temp_directory
454
455    def __exit__(self, type, value, traceback):
456        if self.delete:
457            # Delete temp directory and all of its contents
458            if os.path.isdir(self.temp_directory):
459                delete_directory(self.temp_directory)
460
461
462class TempFilePath:
463    """ Gives a path to a uniquely named temporary file that does not currently exist on __enter__, deletes the file if it exists on __exit__.
464
465    Args:
466        directory (Union[str, None], optional): The directory to create a temporary file in.
467        delete (bool, optional): Default True. Delete the temporary file on on __exit__
468    """
469
470    def __init__(self, directory: Union[str, None] = None, delete: bool = True):
471        # Validate args
472        if not isinstance(directory, (str, type(None))):
473            raise TypeError(directory)
474        if isinstance(directory, str) and not os.path.isdir(directory):
475            raise NotADirectoryError(directory)
476        if not isinstance(delete, bool):
477            raise TypeError(delete)
478
479        self.temp_file: str = None
480        self.directory = directory or tempfile.gettempdir()
481        self.delete = delete
482
483        while self.temp_file is None or os.path.isfile(self.temp_file):
484            self.temp_file = os.path.join(self.directory, next(tempfile._get_candidate_names()))
485
486        # Normalize
487        self.temp_file = str(pathlib.Path(self.temp_file).resolve())
488
489        # Create temp directory if it does not exist
490        os.makedirs(os.path.dirname(self.temp_file), exist_ok=True)
491
492    def __enter__(self):
493        return self.temp_file
494
495    def __exit__(self, type, value, traceback):
496        if self.delete:
497            if os.path.isfile(self.temp_file):
498                os.remove(self.temp_file)
499
500
501# NOTE typehint as string due to no "from __future__ import annotations" support on python 3.6 on ubuntu-16.04 / centos7
502def validate_xml(xml: Union[str, bytes, bytearray, io.BytesIO, "glasswall.content_management.policies.policy.Policy"]):
503    """ Attempts to parse the xml provided, returning the xml as string. Raises ValueError if the xml cannot be parsed.
504
505    Args:
506        xml (Union[str, bytes, bytearray, io.BytesIO, glasswall.content_management.policies.policy.Policy]): The xml string, or file path, bytes, or ContentManagementPolicy instance to parse.
507
508    Returns:
509        xml_string (str): A string representation of the xml.
510
511    Raises:
512        ValueError: if the xml cannot be parsed.
513        TypeError: if the type of arg "xml" is invalid
514    """
515    try:
516        # Get tree from file/str
517        if isinstance(xml, str):
518            try:
519                is_file = os.path.isfile(os.path.abspath(xml))
520            except Exception:
521                is_file = False
522
523            if is_file:
524                tree = etree.parse(xml)
525            else:
526                xml = xml.encode("utf-8")
527                tree = etree.fromstring(xml)
528
529        # Get tree from bytes, bytearray, io.BytesIO
530        elif isinstance(xml, (bytes, bytearray, io.BytesIO)):
531            # Convert bytes, bytearray to io.BytesIO
532            if isinstance(xml, (bytes, bytearray)):
533                xml = as_io_BytesIO(xml)
534            tree = etree.parse(xml)
535
536        # Get tree from ContentManagementPolicy instance
537        elif isinstance(xml, glasswall.content_management.policies.policy.Policy):
538            xml = xml.text.encode("utf-8")
539            tree = etree.fromstring(xml)
540
541        else:
542            raise TypeError(xml)
543
544    except etree.XMLSyntaxError:
545        raise ValueError(xml)
546
547    # # convert tree to string and include xml declaration header utf8
548    etree.indent(tree, space=" " * 4)
549    xml_string = etree.tostring(tree, encoding="utf-8", xml_declaration=True, pretty_print=True).decode()
550
551    return xml_string
552
553
554def xml_as_dict(xml):
555    """ Converts a simple single-level xml into a dictionary.
556
557    Args:
558        xml (Union[str, bytes, bytearray, io.BytesIO]): The xml string, or file path, or bytes to parse.
559
560    Returns:
561        dict_ (dict): A dictionary of element tag : text
562    """
563    # Convert xml to string
564    xml_string = validate_xml(xml)
565
566    # Get root
567    root = etree.fromstring(xml_string.encode())
568
569    dict_ = {
570        element.tag: element.text
571        for element in root
572    }
573
574    # Sort for ease of viewing logs
575    dict_ = {k: v for k, v in sorted(dict_.items())}
576
577    return dict_
578
579
580def deprecated_alias(**aliases: str) -> Callable:
581    """ Decorator for deprecated function and method arguments.
582
583    Use as follows:
584
585    @deprecated_alias(old_arg='new_arg')
586    def myfunc(new_arg):
587        ...
588
589    https://stackoverflow.com/a/49802489
590    """
591
592    def deco(f: Callable):
593        @functools.wraps(f)
594        def wrapper(*args, **kwargs):
595            rename_kwargs(f.__name__, kwargs, aliases)
596            return f(*args, **kwargs)
597
598        return wrapper
599
600    return deco
601
602
603def rename_kwargs(func_name: str, kwargs: Dict[str, Any], aliases: Dict[str, str]):
604    """ Helper function for deprecating function arguments.
605
606    https://stackoverflow.com/a/49802489
607    """
608    for alias, new in aliases.items():
609        if alias in kwargs:
610            if new in kwargs:
611                raise TypeError(
612                    f"{func_name} received both {alias} and {new} as arguments!"
613                    f" {alias} is deprecated, use {new} instead."
614                )
615            warnings.warn(
616                message=(
617                    f"`{alias}` is deprecated as an argument to `{func_name}`; use"
618                    f" `{new}` instead."
619                ),
620                category=DeprecationWarning,
621                stacklevel=3,
622            )
623            kwargs[new] = kwargs.pop(alias)
624
625
626def deprecated_function(replacement_function):
627    def decorator(f: Callable):
628        @functools.wraps(f)
629        def wrapper(*args, **kwargs):
630            warnings.warn(
631                message=f"Call to deprecated method: '{f.__name__}'. Use '{replacement_function.__name__}' instead.",
632                category=DeprecationWarning,
633                stacklevel=3
634            )
635            return replacement_function(*args, **kwargs)
636
637        return wrapper
638
639    return decorator