From e384df00e7b4b1fc131427204adf1b7e5570a73d Mon Sep 17 00:00:00 2001
From: Blake Warner <blake@null3d.com>
Date: Fri, 1 May 2020 13:46:35 -0400
Subject: [PATCH]

---
 Server/Controller/Api.py | 11 +++++++++--
 Server/__init__.py       |  2 +-
 server.py                | 32 +++++++++++++++++++++++++++++---
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/Server/Controller/Api.py b/Server/Controller/Api.py
index 2e7a03763..dc8a99951 100644
--- a/Server/Controller/Api.py
+++ b/Server/Controller/Api.py
@@ -19,6 +19,7 @@
 from googleapiclient.http import MediaIoBaseDownload
 import io
 import hashlib
+import traceback
 
 SCOPES = ['https://www.googleapis.com/auth/drive']
 
@@ -298,6 +299,7 @@ def serveFile(response, path, filename = None, start = None, end = None):
 				status.close()
 	except BaseException as e:
 		Print.error('File download exception: ' + str(e))
+		traceback.print_exc(file=sys.stdout)
 
 	if response.bytesSent == 0:
 		response.write(b'')
@@ -378,6 +380,7 @@ def getDownload(request, response, start = None, end = None):
 				status.close()
 	except BaseException as e:
 		Print.error('NSP download exception: ' + str(e))
+		traceback.print_exc(file=sys.stdout)
 	if response.bytesSent == 0:
 		response.write(b'')
 
@@ -700,6 +703,9 @@ def getGdriveToken(request, response):
 
 		with open('token.pickle', 'wb') as token:
 			pickle.dump(creds, token)
+			
+		with open('gdrive.token', 'w') as token:
+			token.write(json.dumps({'access_token': creds.token, 'refresh_token': creds.refresh_token}))
 	
 	r = {}
 	r['access_token'] = creds.token
@@ -708,8 +714,8 @@ def getGdriveToken(request, response):
 	with open(Config.getGdriveCredentialsFile(), 'r') as f:
 		r['credentials'] = json.loads(f.read())
 	
-	
-	response.write(json.dumps(r))
+	if response is not None:
+		response.write(json.dumps(r))
 		
 def listGdriveDir(path):
 	r = {'dirs': [], 'files': []}
@@ -922,3 +928,4 @@ def getFileSize(request, response):
 		response.write(json.dumps({'success': False, 'message': str(e)}))
 
 
+
diff --git a/Server/__init__.py b/Server/__init__.py
index 3632ea265..8a80bc3b8 100644
--- a/Server/__init__.py
+++ b/Server/__init__.py
@@ -189,7 +189,7 @@ def setMime(self, fileName):
 	def attachFile(self, fileName):
 		#Print.info('Attaching file ' + fileName)
 		self.setMime(fileName)
-		self.headers['Content-Disposition'] = 'attachment; filename=' + fileName
+		self.headers['Content-Disposition'] = 'attachment; filename=' + re.sub(r'[^\x00-\x7F]+','_', fileName)
 
 	def sendHeader(self):
 		self.handler.send_response(self.status)
diff --git a/server.py b/server.py
index 81957673f..df4983c7a 100644
--- a/server.py
+++ b/server.py
@@ -8,7 +8,7 @@
 import urllib3
 import urllib
 import json
-# import webbrowser
+import webbrowser
 import Server
 
 import nut
@@ -28,7 +28,7 @@
 import socket
 
 import sys
-from PyQt5.QtWidgets import QMainWindow, QApplication, QWidget, QAction, QTableWidget,QTableWidgetItem,QVBoxLayout,QDesktopWidget, QTabWidget, QProgressBar, QLabel,QHBoxLayout, QLineEdit, QPushButton, QCheckBox
+from PyQt5.QtWidgets import QMainWindow, QApplication, QWidget, QAction, QTableWidget,QTableWidgetItem,QVBoxLayout,QDesktopWidget, QTabWidget, QProgressBar, QLabel,QHBoxLayout, QLineEdit, QPushButton, QCheckBox, QMessageBox
 from PyQt5.QtGui import QIcon
 from PyQt5.QtCore import pyqtSlot,Qt,QTimer
 from PyQt5 import QtWidgets
@@ -58,6 +58,10 @@ def __init__(self, app):
 		self.scan = QPushButton('Scan', app)
 		self.scan.clicked.connect(app.on_scan)
 		self.layout.addWidget(self.scan)
+		
+		self.gdrive = QPushButton('Setup GDrive OAuth', app)
+		self.gdrive.clicked.connect(app.on_gdrive)
+		self.layout.addWidget(self.gdrive)
 
 		# self.autolaunchBrowser = QCheckBox("Launch Web Browser?", app)
 		# self.autolaunchBrowser.setChecked(Config.autolaunchBrowser)
@@ -144,7 +148,7 @@ def __init__(self):
 		super().__init__()
 		self.setWindowIcon(QIcon('public_html/images/logo.jpg'))
 		screen = QDesktopWidget().screenGeometry()
-		self.title = 'NUT USB / Web Server v2.6'
+		self.title = 'NUT USB / Web Server v2.7'
 		self.left = screen.width() / 4
 		self.top = screen.height() / 4
 		self.width = screen.width() / 2
@@ -202,6 +206,28 @@ def on_scan(self):
 		self.tableWidget.setRowCount(0)
 		nut.scan()
 		self.refreshTable()
+		
+	@pyqtSlot()
+	def on_gdrive(self):
+		if Config.getGdriveCredentialsFile() is None:
+			webbrowser.open_new_tab('https://developers.google.com/drive/api/v3/quickstart/go')
+			QMessageBox.information(self, 'Google Drive OAuth Setup', "You require a credentials.json file to set up Google Drive OAuth.  This file can be obtained from https://developers.google.com/drive/api/v3/quickstart/go , click on the blue button that says 'Enable the Drive API' and save the credentials.json to t his application's directory.")
+		else:
+			buttonReply = QMessageBox.question(self, 'Google Drive OAuth Setup', "Do you you want to setup GDrive OAuth?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
+			
+			if buttonReply == QMessageBox.Yes:
+				try:
+					os.unlink('gdrive.token')
+				except:
+					pass
+					
+				try:
+					os.unlink('token.pickle')
+				except:
+					pass
+					
+				Server.Controller.Api.getGdriveToken(None, None)
+				QMessageBox.information(self, 'Google Drive OAuth Setup', "OAuth has completed.  Please copy gdrive.token and credentials.json to your Nintendo Switch's sdmc:/switch/tinfoil/ and/or sdmc:/switch/sx/ directories.")
 
 	@pyqtSlot()
 	def refreshTable(self):