1 files changed, 56 insertions, 40 deletions
diff --git a/util/drive.go b/util/drive.go
index c4ba137..259912f 100644
--- a/util/drive.go
+++ b/util/drive.go
@@ -25,44 +25,60 @@ func ParentList(parents []*drive.ParentReference) string {
 	return strings.Join(ids, ", ")
 }
 
-func ExportFormat(info *drive.File, format string) (downloadUrl string, extension string, err error) {
-	// See https://developers.google.com/drive/web/manage-downloads#downloading_google_documents
-	switch format {
-	case "docx":
-		extension = ".docx"
-		downloadUrl = info.ExportLinks["application/vnd.openxmlformats-officedocument.wordprocessingml.document"]
-	case "xlsx":
-		extension = ".xlsx"
-		downloadUrl = info.ExportLinks["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"]
-	case "pptx":
-		extension = ".pptx"
-		downloadUrl = info.ExportLinks["application/application/vnd.openxmlformats-officedocument.presentationml.presentation"]
-	case "odf":
-		extension = ".odf"
-		downloadUrl = info.ExportLinks["application/vnd.oasis.opendocument.text"]
-	case "ods":
-		extension = ".ods"
-		downloadUrl = info.ExportLinks["application/x-vnd.oasis.opendocument.spreadsheet"]
-	case "pdf":
-		extension = ".pdf"
-		downloadUrl = info.ExportLinks["application/pdf"]
-	case "rtf":
-		extension = ".rtf"
-		downloadUrl = info.ExportLinks["application/rtf"]
-	case "csv":
-		extension = ".csv"
-		downloadUrl = info.ExportLinks["text/csv"]
-	case "html":
-		extension = ".html"
-		downloadUrl = info.ExportLinks["text/html"]
-	case "txt":
-		extension = ".txt"
-		downloadUrl = info.ExportLinks["text/plain"]
-	case "json":
-		extension = ".json"
-		downloadUrl = info.ExportLinks["application/vnd.google-apps.script+json"]
-	default:
-		err = fmt.Errorf("Unknown export format: %s. Known formats: docx, xlsx, pptx, odf, ods, pdf, rtf, csv, txt, html, json", format)
-	}
-	return
+func InternalDownloadUrlAndExtension(info *drive.File, format string) (downloadUrl string, extension string, err error) {
+    // Make a list of available mime types for this file
+    availableMimeTypes := make([]string, 0)
+    for mime, _ := range info.ExportLinks {
+        availableMimeTypes = append(availableMimeTypes, mime)
+    }
+
+    mimeExtensions := map[string]string{
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx",
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx",
+        "application/application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
+        "application/vnd.oasis.opendocument.text": "odf",
+        "application/x-vnd.oasis.opendocument.spreadsheet": "ods",
+        "application/pdf": "pdf",
+        "application/rtf": "rtf",
+        "text/csv": "csv",
+        "text/html": "html",
+        "text/plain": "txt",
+        "application/vnd.google-apps.script+json": "json",
+    }
+
+    // Make a list of available formats for this file
+    availableFormats :=  make([]string, 0)
+    for _, mime := range availableMimeTypes {
+        if ext, ok := mimeExtensions[mime]; ok {
+            availableFormats = append(availableFormats, ext)
+        }
+    }
+
+    // Return DownloadUrl if no format is specified
+    if format == "" {
+        if info.DownloadUrl == "" {
+            return "", "", fmt.Errorf("A format needs to be specified to download this file (--format). Available formats: %s", strings.Join(availableFormats, ", "))
+        }
+        return info.DownloadUrl, "", nil
+    }
+
+    // Ensure that the specified format is available
+    if !inArray(format, availableFormats) {
+        if len(availableFormats) > 0 {
+            return "", "", fmt.Errorf("Invalid format. Available formats: %s", strings.Join(availableFormats, ", "))
+        } else {
+            return "", "", fmt.Errorf("No export formats are available for this file")
+        }
+    }
+
+    // Grab download url
+    for mime, f := range mimeExtensions {
+        if f == format {
+            downloadUrl = info.ExportLinks[mime]
+            break
+        }
+    }
+
+    extension = "." + format
+    return
 }