Coverage for backpack/autoidentity.py: 83%

69 statements  

« prev     ^ index     » next       coverage.py v7.2.2, created at 2023-03-30 23:12 +0000

1''' This module contains the :class:`AutoIdentity` class that provides information about the 

2application execution environment. ''' 

3 

4import os 

5import logging 

6import datetime 

7from typing import Dict, Optional, Iterator, Any 

8import time 

9 

10import boto3 

11from pydantic import BaseModel, Field 

12 

13class AutoIdentityData(BaseModel): 

14 ''' Data class to store auto identity information. ''' 

15 

16 application_instance_id: str = Field(alias='ApplicationInstanceId') 

17 ''' Application instance id. ''' 

18 

19 application_name: str = Field(alias='Name') 

20 ''' Name of this application. ''' 

21 

22 application_tags: Dict[str, str] = Field(alias='Tags') 

23 ''' Tags associated with this application. ''' 

24 

25 device_id: str = Field(alias='DefaultRuntimeContextDevice') 

26 ''' Device id of the appliance running this application. ''' 

27 

28 device_name: str = Field(alias='DefaultRuntimeContextDeviceName') 

29 ''' Name of this application. ''' 

30 

31 application_created_time: datetime.datetime = Field(alias='CreatedTime') 

32 ''' Application deployment time. ''' 

33 

34 application_status: str = Field(alias='HealthStatus') 

35 ''' Health status of this application. ''' 

36 

37 application_description: Optional[str] = Field(alias='Description') 

38 ''' The description of this application. ''' 

39 

40 @classmethod 

41 def for_test_environment(cls, application_instance_id: str, application_name: str): 

42 ''' Initializes a dummy AutoIdentityData to be used in test environment. ''' 

43 return cls( 

44 ApplicationInstanceId=application_instance_id, 

45 Name=application_name, 

46 Tags={}, 

47 DefaultRuntimeContextDevice='emulator', 

48 DefaultRuntimeContextDeviceName='test_utility_emulator', 

49 CreatedTime=datetime.datetime.now(), 

50 HealthStatus='TEST_UTILITY', 

51 Description=application_name, 

52 ) 

53 

54 

55class AutoIdentityError(RuntimeError): 

56 ''' AutoIdentity specific error. ''' 

57 

58 

59class AutoIdentityFetcher: 

60 ''' AutoIdentity instance queries metadata of the current application instance. 

61 

62 The IAM policy associated with the `Panorama Application Role`_ 

63 of this app should grant the execution of 

64 `panorama:ListApplicationInstances`_ operation. 

65 

66 Args: 

67 device_region: The AWS region where this Panorama appliance is registered. 

68 application_instance_id: The application instance id. If left to `None`, 

69 :class:`AutoIdentity` will try to find the instance id in the environment variable. 

70 parent_logger: If you want to connect the logger to a parent, specify it here. 

71 

72 .. _`Panorama Application Role`: 

73 https://docs.aws.amazon.com/panorama/latest/dev/permissions-application.html 

74 .. _`panorama:ListApplicationInstances`: 

75 https://docs.aws.amazon.com/service-authorization/latest/reference/list_awspanorama.html#awspanorama-actions-as-permissions 

76 ''' 

77 

78 # pylint: disable=too-many-instance-attributes,too-few-public-methods 

79 # This class functions as a data class that reads its values from the environment 

80 

81 def __init__(self, 

82 device_region: str, 

83 application_instance_id: Optional[str] = None, 

84 parent_logger: Optional[logging.Logger] = None 

85 ): 

86 self._logger = ( 

87 logging.getLogger(self.__class__.__name__) if parent_logger is None else 

88 parent_logger.getChild(self.__class__.__name__) 

89 ) 

90 self.application_instance_id = ( 

91 application_instance_id or os.environ.get('AppGraph_Uid') 

92 ) 

93 if not self.application_instance_id: 

94 raise AutoIdentityError( 

95 'Could not find application instance id in environment variable "AppGraph_Uid"' 

96 ) 

97 self.device_region = device_region 

98 

99 def get_data(self, 

100 retry_freq: Optional[float] = 5, 

101 max_retry_freq: Optional[float] = 300, 

102 max_retry_num: Optional[int] = 50) -> AutoIdentityData: 

103 ''' Fetches the auto identity data. 

104 

105 Args: 

106 retry_freq (Optional[float]): If set to a float number, AutoIdentity will keep retrying 

107 fetching the auto identity data from remote services if the status of the app was 

108 "NOT_AVAILABLE". If set to None, will not retry. 

109 

110 Raises: 

111 AutoIdentityError: if could not fetch the auto identity information, and retry_freq is set 

112 to None. 

113 ''' 

114 

115 def fetch() -> Dict[str, Any]: 

116 app_instance_data = self._app_instance_data(self.application_instance_id) 

117 self._logger.info('Fetched data: %s', app_instance_data) 

118 if not app_instance_data: 

119 raise AutoIdentityError( 

120 'Could not find application instance in service response. ' 

121 'Check if application_instance_id={} ' 

122 'and device_region={} parameters are correct.' 

123 .format(self.application_instance_id, self.device_region) 

124 ) 

125 else: 

126 return app_instance_data 

127 

128 retries = 0 

129 while True: 

130 app_instance_data = fetch() 

131 status = app_instance_data.get('HealthStatus', 'NOT_AVAILABLE') 

132 self._logger.info('Application HealthStatus=%s', status) 

133 if status in ('NOT_AVAILABLE'): 

134 if retry_freq is None: 

135 raise AutoIdentityError( 

136 f'Application HealthStatus is "{status}" and retry is disabled.' 

137 ) 

138 else: 

139 self._logger.info('Will retry fetching auto identity data in %s seconds.', retry_freq) 

140 time.sleep(retry_freq) 

141 retry_freq = retry_freq * 2 

142 if max_retry_freq is not None: 

143 retry_freq = min(retry_freq, max_retry_freq) 

144 retries += 1 

145 if max_retry_num is not None and retries > max_retry_num: 

146 raise AutoIdentityError('Maximum number of retries reached.') 

147 else: 

148 continue 

149 else: 

150 return AutoIdentityData(**app_instance_data) 

151 

152 def _list_app_instances(self, deployed_only=True) -> Iterator[Dict[str, Any]]: 

153 session = boto3.Session(region_name=self.device_region) 

154 panorama = session.client('panorama') 

155 next_token = None 

156 while True: 

157 kwargs = {'StatusFilter': 'DEPLOYMENT_SUCCEEDED'} if deployed_only else {} 

158 if next_token: 

159 kwargs['NextToken'] = next_token 

160 response = panorama.list_application_instances(**kwargs) 

161 inst: Dict[str, Any] 

162 for inst in response['ApplicationInstances']: 

163 yield inst 

164 if 'NextToken' in response: 

165 next_token = response['NextToken'] 

166 else: 

167 break 

168 

169 def _app_instance_data(self, application_instance_id) -> Optional[Dict[str, Any]]: 

170 matches = (inst 

171 for inst in self._list_app_instances() 

172 if inst.get('ApplicationInstanceId') == application_instance_id 

173 ) 

174 return next(matches, None)