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
« 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 """
4import json
5import os
7from ..utils.logger_config import setup_logger, DEFAULT_LOGGING_LEVEL
8from ..utils.reporter_utils import delete_file
11class GoogleDriveUploader:
12 """Google Drive uploader class"""
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
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
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
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)
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"]
79 def __get_new_oauth_token(self, device_code):
80 """
81 Get new OAuth token
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"]
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
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
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
125 :return: Google OAuth access token, refresh toke (strings)
126 """
127 device_code, user_code, url = self.__get_new_device_codes()
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")
139 access_token, refresh_token = self.__get_new_oauth_token(device_code=device_code)
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 )
148 return access_token, refresh_token
150 def __upload_to_gdrive(self, filename=None, mime_type=None):
151 """
152 Upload file to Google Drive using access token
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")
177 # Flow
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
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)
194 def upload(self, filename=None, mime_type=None):
195 """
196 Upload file to Google Drive and cleanup, if needed.
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 )