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

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

2""" Google Drive uploader module """

3

4 import json

5 import os

6

7 from ..utils.logger_config import setup_logger, DEFAULT_LOGGING_LEVEL

8 from ..utils.reporter_utils import delete_file

9

10

11 class 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 )