From f7f869736b27afadaaaf6c8a754a7b29da2749c6 Mon Sep 17 00:00:00 2001 From: Jeroen Oudshoorn Date: Sat, 14 Oct 2023 23:21:27 +0200 Subject: [PATCH] v2.4.7 Signed-off-by: Jeroen Oudshoorn --- pwnagotchi/plugins/default/gdrivesync.py | 133 ++++++++++++++--------- 1 file changed, 83 insertions(+), 50 deletions(-) diff --git a/pwnagotchi/plugins/default/gdrivesync.py b/pwnagotchi/plugins/default/gdrivesync.py index 2ad3e8b3..60929afb 100644 --- a/pwnagotchi/plugins/default/gdrivesync.py +++ b/pwnagotchi/plugins/default/gdrivesync.py @@ -63,19 +63,21 @@ class GdriveSync(plugins.Plugin): # if backup file does not exist, we will check for backup folder on gdrive. if not self.backup: # Use self.options['backup_folder'] as the folder ID where backups are stored - backup_folder_id = self.create_folder_if_not_exists(self.options['backup_folder']) + backup_folder, root_folder, pwnagotchi_folder = self.create_folder_if_not_exists(self.options['backup_folder']) # Continue with the rest of the code using backup_folder_id - file_list = self.drive.ListFile({'q': f"'{backup_folder_id}' in parents and mimeType = 'application/vnd.google-apps.folder' and trashed=false"}).GetList() + root_file_list = self.drive.ListFile({ + 'q': f"'{root_folder}' in parents and mimeType = 'application/vnd.google-apps.folder' and trashed=false"}).GetList() + pwnagotchi_file_list = self.drive.ListFile({'q': f"'{pwnagotchi_folder}' in parents and mimeType = 'application/vnd.google-apps.folder' and trashed=false"}).GetList() - if not file_list: + if not (root_file_list or pwnagotchi_file_list): # Handle the case where no files were found - logging.warning(f"[gDriveSync] No files found in the folder with ID {backup_folder_id}") + # logging.warning(f"[gDriveSync] No files found in the folder with ID {root_file_list} and {pwnagotchi_file_list}") if self.options['backupfiles'] is not None: self.backupfiles = self.backupfiles + self.options['backupfiles'] self.backup_files(self.backupfiles, '/backup') - self.upload_to_gdrive('/backup', self.options['backup_folder']) + self.upload_to_gdrive('/backup', self.options['backup_folder'],backup_folder, root_folder, pwnagotchi_folder) self.backup = True # Specify the local backup path @@ -84,11 +86,17 @@ class GdriveSync(plugins.Plugin): # Create the local backup directory if it doesn't exist os.makedirs(local_backup_path, exist_ok=True) - # Download each file in the folder - for file in file_list: - local_file_path = os.path.join(local_backup_path, file['title']) - file.GetContentFile(local_file_path) - logging.info(f"[gDriveSync] Downloaded {file['title']} from Google Drive") + # Download each file in the /root folder + for root_file in root_file_list: + local_file_path = os.path.join(local_backup_path, root_file['title']) + root_file.GetContentFile(local_file_path) + logging.info(f"[gDriveSync] Downloaded {root_file['title']} from Google Drive") + + # Download each file in the /etc/pwnagotchi folder + for pwnagotchi_file in pwnagotchi_file_list: + local_file_path = os.path.join(local_backup_path, pwnagotchi_file['title']) + pwnagotchi_file.GetContentFile(local_file_path) + logging.info(f"[gDriveSync] Downloaded {pwnagotchi_file['title']} from Google Drive") # Optionally, you can use the downloaded files as needed # For example, you can copy them to the corresponding directories @@ -104,49 +112,65 @@ class GdriveSync(plugins.Plugin): logging.error(f"Error: {e}") self.ready = False - def get_folder_id_by_name(self, drive, folder_name): - file_list = drive.ListFile({'q': "mimeType='application/vnd.google-apps.folder' and trashed=false"}).GetList() + def get_folder_id_by_name(self, drive, folder_name, parent_folder_id=None): + query = "mimeType='application/vnd.google-apps.folder' and trashed=false" + if parent_folder_id: + query += f" and '{parent_folder_id}' in parents" + + file_list = drive.ListFile({'q': query}).GetList() for file in file_list: if file['title'] == folder_name: return file['id'] - return None + return None def create_folder_if_not_exists(self, backup_folder_name): - # First, create or retrieve the *BACKUP_FOLDER* folder + # First, try to retrieve the existing *BACKUP_FOLDER* folder backup_folder_id = self.get_folder_id_by_name(self.drive, backup_folder_name) + if backup_folder_id is None: - # Now create /*BACKUP_FOLDER* + # If not found, create *BACKUP_FOLDER* backup_folder = self.drive.CreateFile( - {'title': 'PwnagotchiBackups', 'mimeType': 'application/vnd.google-apps.folder'}) + {'title': backup_folder_name, 'mimeType': 'application/vnd.google-apps.folder'}) backup_folder.Upload() backup_folder_id = backup_folder['id'] - logging.info(f"[gDriveSync] Created folder 'PwnagotchiBackups' with ID: {backup_folder_id}") + logging.info(f"[gDriveSync] Created folder '{backup_folder_name}' with ID: {backup_folder_id}") - # Now, create /*BACKUP_FOLDER*/root - root_folder = self.drive.CreateFile( - {'title': 'root', 'mimeType': 'application/vnd.google-apps.folder', - 'parents': [{'id': backup_folder_id}]}) - root_folder.Upload() - root_folder_id = root_folder['id'] - logging.info(f"[gDriveSync] Created folder 'root' with ID: {root_folder_id}") + # Now, try to retrieve or create /*BACKUP_FOLDER*/root + root_folder_id = self.get_or_create_subfolder('root', backup_folder_id) - # Now, create /*BACKUP_FOLDER*/etc - etc_folder = self.drive.CreateFile( - {'title': 'etc', 'mimeType': 'application/vnd.google-apps.folder', - 'parents': [{'id': backup_folder_id}]}) - etc_folder.Upload() - etc_folder_id = etc_folder['id'] - logging.info(f"[gDriveSync] Created folder 'etc' with ID: {etc_folder_id}") + # Now, try to retrieve or create /*BACKUP_FOLDER*/etc + etc_folder_id = self.get_or_create_subfolder('etc', backup_folder_id) - # Now, create /*BACKUP_FOLDER*/etc/pwnagotchi - pwnagotchi_folder = self.drive.CreateFile( - {'title': 'pwnagotchi', 'mimeType': 'application/vnd.google-apps.folder', - 'parents': [{'id': etc_folder_id}]}) - pwnagotchi_folder.Upload() - pwnagotchi_folder_id = pwnagotchi_folder['id'] - logging.info(f"[gDriveSync] Created folder 'pwnagotchi' under 'etc' with ID: {pwnagotchi_folder_id}") + # Now, try to retrieve or create /*BACKUP_FOLDER*/etc/pwnagotchi + pwnagotchi_folder_id = self.get_or_create_subfolder('pwnagotchi', etc_folder_id) - return root_folder_id, pwnagotchi_folder_id # Return the IDs of both root and pwnagotchi folders + return backup_folder_id, root_folder_id, pwnagotchi_folder_id # Return the IDs of both root and pwnagotchi folders + else: + # If found, also try to retrieve or create /*BACKUP_FOLDER*/root + root_folder_id = self.get_or_create_subfolder('root', backup_folder_id) + + # Also, try to retrieve or create /*BACKUP_FOLDER*/etc + etc_folder_id = self.get_or_create_subfolder('etc', backup_folder_id) + + # Also, try to retrieve or create /*BACKUP_FOLDER*/etc/pwnagotchi + pwnagotchi_folder_id = self.get_or_create_subfolder('pwnagotchi', etc_folder_id) + + return backup_folder_id, root_folder_id, pwnagotchi_folder_id # Return the IDs of both root and pwnagotchi folders + + def get_or_create_subfolder(self, subfolder_name, parent_folder_id): + # Try to retrieve the subfolder + subfolder_id = self.get_folder_id_by_name(self.drive, subfolder_name, parent_folder_id) + + if subfolder_id is None: + # If not found, create the subfolder + subfolder = self.drive.CreateFile( + {'title': subfolder_name, 'mimeType': 'application/vnd.google-apps.folder', + 'parents': [{'id': parent_folder_id}]}) + subfolder.Upload() + subfolder_id = subfolder['id'] + logging.info(f"[gDriveSync] Created folder '{subfolder_name}' with ID: {subfolder_id}") + + return subfolder_id def on_unload(self, ui): logging.info("[gdrivesync] unloaded") @@ -192,19 +216,28 @@ class GdriveSync(plugins.Plugin): except Exception as e: logging.error(f"[gDriveSync] Error during backup_path: {e}") - def upload_to_gdrive(self, backup_path, gdrive_folder): + def upload_to_gdrive(self, backup_path, gdrive_folder, root_folder_id, pwnagotchi_folder_id): try: - # Get the ID of the main backup folder - backup_folder_id = self.get_folder_id_by_name(self.drive, gdrive_folder) - folder = self.drive.CreateFile({'id': backup_folder_id}) + # Upload files to the /root folder + if root_folder_id is not None: + root_folder = self.drive.CreateFile({'id': root_folder_id}) + for root, dirs, files in os.walk('/root'): + for filename in files: + file_path = os.path.join(root, filename) + gdrive_file = self.drive.CreateFile({'title': filename, 'parents': [{'id': root_folder['id']}]}) + gdrive_file.Upload() + logging.info(f"[gDriveSync] Uploaded {file_path} to Google Drive") - # Upload files and folders to the created folder - for root, dirs, files in os.walk(backup_path): - for filename in files: - file_path = os.path.join(root, filename) - gdrive_file = self.drive.CreateFile({'title': filename, 'parents': [{'id': folder['id']}]}) - gdrive_file.Upload() - logging.info(f"[gDriveSync] Uploaded {file_path} to Google Drive") + # Upload files to the /etc/pwnagotchi folder + if pwnagotchi_folder_id is not None: + pwnagotchi_folder = self.drive.CreateFile({'id': pwnagotchi_folder_id}) + for root, dirs, files in os.walk('/etc/pwnagotchi'): + for filename in files: + file_path = os.path.join(root, filename) + gdrive_file = self.drive.CreateFile( + {'title': filename, 'parents': [{'id': pwnagotchi_folder['id']}]}) + gdrive_file.Upload() + logging.info(f"[gDriveSync] Uploaded {file_path} to Google Drive") logging.info(f"[gDriveSync] Backup uploaded to Google Drive") except pydrive2.files.ApiRequestError as api_error: