1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6"""
7Utility decorators
8
9This module will be executed in separate process after az process is terminated to upload traces, so it is preferable
10that it doesn't import modules other than those in the Python Standard Library
11"""
12
13import hashlib
14from functools import wraps
15
16from knack.log import get_logger
17
18
19# pylint: disable=too-few-public-methods
20class Completer:
21
22    def __init__(self, func):
23        self.func = func
24
25    def __call__(self, **kwargs):
26        namespace = kwargs['parsed_args']
27        prefix = kwargs['prefix']
28        cmd = namespace._cmd  # pylint: disable=protected-access
29        return self.func(cmd, prefix, namespace)
30
31
32def call_once(factory_func):
33    """"
34    When a function is annotated by this decorator, it will be only executed once. The result will be cached and
35    returned for following invocations.
36    """
37    factory_func.executed = False
38    factory_func.cached_result = None
39
40    def _wrapped(*args, **kwargs):
41        if not factory_func.executed:
42            factory_func.cached_result = factory_func(*args, **kwargs)
43
44        return factory_func.cached_result
45
46    return _wrapped
47
48
49def hash256_result(func):
50    """
51    Secure the return string of the annotated function with SHA256 algorithm. If the annotated function doesn't return
52    string or return None, raise ValueError.
53    """
54
55    @wraps(func)
56    def _decorator(*args, **kwargs):
57        val = func(*args, **kwargs)
58        if val is None:
59            raise ValueError('Return value is None')
60        if not isinstance(val, str):
61            raise ValueError('Return value is not string')
62        if not val:
63            return val
64        hash_object = hashlib.sha256(val.encode('utf-8'))
65        return str(hash_object.hexdigest())
66
67    return _decorator
68
69
70def suppress_all_exceptions(fallback_return=None, **kwargs):  # pylint: disable=unused-argument
71    # The kwargs is a fallback to ensure extensions (eg. alias) are not broken
72    def _decorator(func):
73        @wraps(func)
74        def _wrapped_func(*args, **kwargs):
75            try:
76                return func(*args, **kwargs)
77            except Exception:  # nopa pylint: disable=broad-except
78                import traceback
79                get_logger(__name__).info('Suppress exception:\n%s', traceback.format_exc())
80                if fallback_return is not None:
81                    return fallback_return
82        return _wrapped_func
83    return _decorator
84