Coverage for testrail_api_reporter/publishers/gdrive_uploader.py: 19%

68 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-08-29 15:21 +0200

1# -*- coding: utf-8 -*- 

2""" Google Drive uploader module """ 

3 

4import json 

5import os 

6 

7from ..utils.logger_config import setup_logger, DEFAULT_LOGGING_LEVEL 

8from ..utils.reporter_utils import delete_file 

9 

10 

11class GoogleDriveUploader: 

12 """Google Drive uploader class""" 

13 

14 # Google token needs to be configured firstly, to do it, you have to visit: 

15 # https://console.developers.google.com/apis/credentials?pli=1 

16 # Create Credentials => OAuth client ID => TV and limited Input Devices and get client_id and a client_secret 

17 # Then pass it as google_id = client_id and google_secret = client_secret 

18 def __init__( 

19 self, 

20 google_id, 

21 google_secret, 

22 google_api_refresh_token=None, 

23 cleanup_needed=True, 

24 backup_filename="backup.zip", 

25 mime_type="application/zip", 

26 logger=None, 

27 log_level=DEFAULT_LOGGING_LEVEL, 

28 ): 

29 """ 

30 General init 

31 

32 :param google_id: Google OAuth client_id, string, required 

33 :param google_secret: Google OAuth client_secret, string, required 

34 :param google_api_refresh_token: Google OAuth refresh token, string, optional 

35 :param cleanup_needed: delete or not backup file after upload, bool, True or False, by default is True 

36 :param backup_filename: custom backup filename, which will be uploaded to GDrive, string, optional 

37 :param mime_type: MIME type of file for upload, string, by default is 'application/zip' 

38 :param logger: logger object, optional 

39 :param log_level: logging level, optional, by default is 'logging.DEBUG' 

40 """ 

41 if not logger: 

42 self.___logger = setup_logger( 

43 name="GoogleDriveUploader", log_file="GoogleDriveUploader.log", level=log_level 

44 ) 

45 else: 

46 self.___logger = logger 

47 self.___logger.debug("Initializing Google Drive Uploader") 

48 # Google 

49 self.__g_id = google_id 

50 self.__g_secret = google_secret 

51 self.__g_token = None 

52 # Service 

53 self.__cleanup_needed = cleanup_needed 

54 self.__backup_filename = backup_filename 

55 self.__mime_type = mime_type 

56 

57 if not google_api_refresh_token or google_api_refresh_token == "": 

58 self.__g_token, self.__g_refresh_token = self.__first_run() 

59 else: 

60 self.__g_refresh_token = google_api_refresh_token 

61 

62 def __get_new_device_codes(self): 

63 """ 

64 Get OAuth codes from Google Drive for a new device (device code and one-time user code) 

65 

66 :return: device_code, user_code, verification_url (strings) 

67 """ 

68 self.___logger.debug("Get temporary Device ID and user code from Google Auth engine") 

69 response = json.loads( 

70 os.popen( 

71 f"curl " 

72 f'-d "client_id={self.__g_id}&scope=https://www.googleapis.com/auth/drive.file"' 

73 f" https://oauth2.googleapis.com/device/code" 

74 ).read() 

75 ) 

76 self.___logger.debug("Response from Google Auth engine: %s", response) 

77 return response["device_code"], response["user_code"], response["verification_url"] 

78 

79 def __get_new_oauth_token(self, device_code): 

80 """ 

81 Get new OAuth token 

82 

83 :param device_code: unique device_code generated by Google OAuth API, String, required 

84 :return: Google OAuth access token, refresh toke (strings) 

85 """ 

86 self.___logger.debug("Get OAuth token from google using Device ID") 

87 response = json.loads( 

88 os.popen( 

89 f"curl -d client_id={self.__g_id} -d client_secret={self.__g_secret} " 

90 f"-d device_code={device_code} " 

91 f"-d grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code " 

92 f"https://accounts.google.com/o/oauth2/token" 

93 ).read() 

94 ) 

95 return response["access_token"], response["refresh_token"] 

96 

97 def __refresh_token(self): 

98 """ 

99 Refresh Google OAuth token 

100 Google token has a limited lifetime (1 hour), so, it needs to be updated from time-to-time 

101 

102 :return: Google OAuth access token (string) 

103 """ 

104 self.___logger.debug("Google OAuth token needs to be refreshed, so, let's do this") 

105 response = json.loads( 

106 os.popen( 

107 f"curl -d client_id={self.__g_id} -d client_secret={self.__g_secret} " 

108 f"-d refresh_token={self.__g_refresh_token} " 

109 f"-d grant_type=refresh_token " 

110 f"https://accounts.google.com/o/oauth2/token" 

111 ).read() 

112 ) 

113 self.__g_token = response["access_token"] 

114 return self.__g_token 

115 

116 def __first_run(self): 

117 """ 

118 In case when user did not provide refresh_token, new access and refresh tokens should be obtained. 

119 To do that, need to: 

120 1) Generate new device code and get user_code confirmation 

121 2) Fill this code to a Google account (i.e., via web browser) 

122 3) Activate API access token permissions 

123 4) Generate new access tokens and refresh tokens 

124 

125 :return: Google OAuth access token, refresh toke (strings) 

126 """ 

127 device_code, user_code, url = self.__get_new_device_codes() 

128 

129 print(f"Please fill device code {user_code} into web browser URL: {url}") 

130 input("When your code will be submitted and account, press enter") 

131 print( 

132 "Now you must ensure that your access rights are granted for this device! Proceed to:\n" 

133 "https://console.developers.google.com/apis/api/drive.googleapis.com/overview\n" 

134 "and open Credentials tab, now confirm OAuth API permissions for this device.\n" 

135 "After submit please wait at least 5 minutes." 

136 ) 

137 input("When 5 minutes passed, press any enter") 

138 

139 access_token, refresh_token = self.__get_new_oauth_token(device_code=device_code) 

140 

141 print( 

142 f"Your access token is:\n{access_token}\nYour refresh token is:\n{refresh_token}\n" 

143 "Please save these credentials secure!\nYour access token will be valid for an 1 hour. " 

144 "If you plan use it in advance, you need to refresh it every hour or before use any time. \n" 

145 "Next time init this class with your refresh token to update access token automatically." 

146 ) 

147 

148 return access_token, refresh_token 

149 

150 def __upload_to_gdrive(self, filename=None, mime_type=None): 

151 """ 

152 Upload file to Google Drive using access token 

153 

154 :param filename: filename to upload, string 

155 :param mime_type: MIME type of file, string 

156 """ 

157 if not filename: 

158 filename = self.__backup_filename 

159 if not mime_type: 

160 mime_type = self.__mime_type 

161 self.___logger.debug("Uploading %s to GoogleDrive", filename) 

162 response = json.loads( 

163 os.popen( 

164 f'curl -X POST -L -H "Authorization: Bearer {self.__g_token}" ' 

165 f'-F "metadata={ name :\'{filename.split(".")[0]}\'} ;' 

166 f'type=application/json;charset=UTF-8" ' 

167 f'-F "file=@{filename};type={mime_type}" ' 

168 f'"https://www.googleapis.com/upload/drive/v3/' 

169 f'files?uploadType=multipart"' 

170 ).read() 

171 ) 

172 if response["id"]: 

173 self.___logger.debug("Backup archive %s was uploaded to Google Drive", filename) 

174 else: 

175 self.___logger.error("Something wrong, please check backup manually or re-run") 

176 

177 # Flow 

178 

179 def __proceed_upload(self, filename=None, mime_type=None): 

180 """ 

181 Prepare valid token (update if needed), then upload the file to Google Drive using access token 

182 

183 :param filename: filename to upload, string 

184 :param mime_type: MIME type of file, string 

185 """ 

186 if not filename: 

187 filename = self.__backup_filename 

188 if not mime_type: 

189 mime_type = self.__mime_type 

190 if not self.__g_token: 

191 self.__refresh_token() 

192 self.__upload_to_gdrive(filename=filename, mime_type=mime_type) 

193 

194 def upload(self, filename=None, mime_type=None): 

195 """ 

196 Upload file to Google Drive and cleanup, if needed. 

197 

198 :param filename: filename to upload, string 

199 :param mime_type: MIME type of file, string 

200 """ 

201 if not filename: 

202 filename = self.__backup_filename 

203 if not mime_type: 

204 mime_type = self.__mime_type 

205 self.__proceed_upload(filename=filename, mime_type=mime_type) 

206 if self.__cleanup_needed: 

207 delete_file( 

208 filename=filename, 

209 debug=self.___logger.debug == DEFAULT_LOGGING_LEVEL, 

210 logger=self.___logger, 

211 )