Add support for uploading files and other media types

1 files changed, 191 insertions, 61 deletions

diff --git a/ui/commands.go b/ui/commands.go
index 6e9d3b7..3636c5b 100644
--- a/ui/commands.go
+++ b/ui/commands.go
@@ -24,6 +24,7 @@ import (
 	"io"
 	"io/ioutil"
 	"math"
+	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -181,37 +182,69 @@ func cmdDownload(cmd *Command) {
 	cmd.Room.StartSelecting(SelectDownload, strings.Join(cmd.Args, " "))
 }
 
-type VideoInfo struct {
+type AVStream struct {
+	CodecType string  `json:"codec_type"`
+	Width     float64 `json:"width"`
+	Height    float64 `json:"height"`
+}
+
+type AVFormat struct {
+	Duration string `json:"duration"`
+}
+
+type AVMediaInfo struct {
+	Streams []AVStream `json:"streams"`
+	Format  AVFormat   `json:"format"`
+}
+
+type AVInfo struct {
 	width    int
 	height   int
 	duration int
+	isVideo  bool
 }
 
-func videoGetInfo(filePath string) (VideoInfo, error) {
+func avGetInfo(filePath string) (AVInfo, error) {
 	cmd := exec.Command("ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", "-show_entries", "format=duration", "--", filePath)
 	var buffer bytes.Buffer
 	cmd.Stdout = &buffer
 	if cmd.Run() != nil {
-		return VideoInfo{0, 0, 0}, errors.New("Failed to extract video dimensions using ffprobe")
+		return AVInfo{0, 0, 0, false}, errors.New("Failed to extract video dimensions using ffprobe")
 	}
 
-	var jsonData map[string]interface{}
+	var jsonData AVMediaInfo
 	if err := json.Unmarshal(buffer.Bytes(), &jsonData); err != nil {
-		return VideoInfo{0, 0, 0}, errors.New("Failed to parse ffprobe response as json")
+		return AVInfo{0, 0, 0, false}, errors.New("Failed to parse ffprobe response as json")
 	}
 
-	// TODO: Check if fields do not exist or a are wrong type???
-	format := jsonData["format"].(map[string]interface{})
-	durationStr := format["duration"].(string)
-	duration, err := strconv.ParseFloat(durationStr, 64)
+	duration, err := strconv.ParseFloat(jsonData.Format.Duration, 64)
 	if err != nil {
-		return VideoInfo{0, 0, 0}, errors.New("Failed to parse ffprobe response as json (duration is not a string)")
+		return AVInfo{0, 0, 0, false}, errors.New("Failed to parse ffprobe response as json (duration is not a string)")
 	}
 
-	streams := jsonData["streams"].([]interface{})
-	firstStream := streams[0].(map[string]interface{})
+	videoStreamIndex := -1
+	audioStreamIndex := -1
+	for i, stream := range jsonData.Streams {
+		if stream.CodecType == "video" {
+			if videoStreamIndex == -1 {
+				videoStreamIndex = i
+			}
+		} else if stream.CodecType == "audio" {
+			if audioStreamIndex == -1 {
+				audioStreamIndex = i
+			}
+		}
+	}
 
-	return VideoInfo{int(firstStream["width"].(float64)), int(firstStream["height"].(float64)), int(duration * 1000)}, nil
+	if videoStreamIndex != -1 {
+		videoStream := jsonData.Streams[videoStreamIndex]
+		// TODO: float64 to int can fail (width, height and duration)
+		return AVInfo{int(videoStream.Width), int(videoStream.Height), int(duration * 1000), true}, nil
+	} else if audioStreamIndex != -1 {
+		return AVInfo{0, 0, int(duration * 1000), false}, nil
+	} else {
+		return AVInfo{0, 0, 0, false}, errors.New("Media file is missing video and audio streams")
+	}
 }
 
 func videoGetThumbnail(filePath string) ([]byte, error) {
@@ -224,29 +257,47 @@ func videoGetThumbnail(filePath string) ([]byte, error) {
 	return buffer.Bytes(), nil
 }
 
-func uploadVideo(cmd *Command, filePath string, fileName string, fileSize int, videoFormat string) {
-	videoInfo, err := videoGetInfo(filePath)
+func uploadAudioVideo(cmd *Command, filePath string, fileName string, fileSize int, contentType string) {
+	avInfo, err := avGetInfo(filePath)
 	if err != nil {
-		cmd.Reply("Failed to get video info, error: %v", err)
+		cmd.Reply("Failed to get audio/video info, error: %v", err)
 		return
 	}
 
-	thumbnailData, err := videoGetThumbnail(filePath)
-	if err != nil {
-		cmd.Reply("Failed to extract thumbnail for video")
-		return
-	}
+	var fileInfo event.FileInfo
+	if avInfo.isVideo {
+		thumbnailData, err := videoGetThumbnail(filePath)
+		if err != nil {
+			cmd.Reply("Failed to extract thumbnail for video")
+			return
+		}
 
-	imageMetadata, _, err := image.DecodeConfig(bytes.NewReader(thumbnailData))
-	if err != nil {
-		cmd.Reply("Failed to decode thumbnail image")
-		return
-	}
+		imageMetadata, _, err := image.DecodeConfig(bytes.NewReader(thumbnailData))
+		if err != nil {
+			cmd.Reply("Failed to decode thumbnail image")
+			return
+		}
 
-	thumbnailResp, err := cmd.Matrix.Client().UploadBytesWithName(thumbnailData, "image/jpeg", "thumbnail.jpg")
-	if err != nil {
-		cmd.Reply("Failed to upload thumbnail: %v", err)
-		return
+		thumbnailResp, err := cmd.Matrix.Client().UploadBytesWithName(thumbnailData, "image/jpeg", "thumbnail.jpg")
+		if err != nil {
+			cmd.Reply("Failed to upload thumbnail: %v", err)
+			return
+		}
+
+		fileInfo = event.FileInfo{
+			Duration:     avInfo.duration,
+			Width:        avInfo.width,
+			Height:       avInfo.height,
+			MimeType:     contentType,
+			Size:         fileSize,
+			ThumbnailURL: id.ContentURIString(fmt.Sprintf("mxc://%s/%s", thumbnailResp.ContentURI.Homeserver, thumbnailResp.ContentURI.FileID)),
+			ThumbnailInfo: &event.FileInfo{
+				Width:    imageMetadata.Width,
+				Height:   imageMetadata.Height,
+				MimeType: "image/jpeg",
+				Size:     len(thumbnailData),
+			},
+		}
 	}
 
 	fileData, err := ioutil.ReadFile(filePath)
@@ -255,32 +306,32 @@ func uploadVideo(cmd *Command, filePath string, fileName string, fileSize int, v
 		return
 	}
 
-	videoMimeType := fmt.Sprintf("video/%s", videoFormat)
-	resp, err := cmd.Matrix.Client().UploadBytesWithName(fileData, videoMimeType, fileName)
+	resp, err := cmd.Matrix.Client().UploadBytesWithName(fileData, contentType, fileName)
 	if err != nil {
-		cmd.Reply("Failed to upload video: %v", err)
+		cmd.Reply("Failed to upload audio/video: %v", err)
 		return
 	}
 
 	txnID := cmd.Matrix.Client().TxnID()
-	content := event.MessageEventContent{
-		Body: fileName,
-		Info: &event.FileInfo{
-			Duration:     videoInfo.duration,
-			Width:        videoInfo.width,
-			Height:       videoInfo.height,
-			MimeType:     videoMimeType,
-			Size:         fileSize,
-			ThumbnailURL: id.ContentURIString(fmt.Sprintf("mxc://%s/%s", thumbnailResp.ContentURI.Homeserver, thumbnailResp.ContentURI.FileID)),
-			ThumbnailInfo: &event.FileInfo{
-				Width:    imageMetadata.Width,
-				Height:   imageMetadata.Height,
-				MimeType: "image/jpeg",
-				Size:     len(thumbnailData),
+	var content event.MessageEventContent
+	if avInfo.isVideo {
+		content = event.MessageEventContent{
+			Body:    fileName,
+			Info:    &fileInfo,
+			MsgType: event.MsgVideo,
+			URL:     id.ContentURIString(fmt.Sprintf("mxc://%s/%s", resp.ContentURI.Homeserver, resp.ContentURI.FileID)),
+		}
+	} else {
+		content = event.MessageEventContent{
+			Body: fileName,
+			Info: &event.FileInfo{
+				Duration: avInfo.duration,
+				MimeType: contentType,
+				Size:     fileSize,
 			},
-		},
-		MsgType: event.MsgVideo,
-		URL:     id.ContentURIString(fmt.Sprintf("mxc://%s/%s", resp.ContentURI.Homeserver, resp.ContentURI.FileID)),
+			MsgType: event.MsgAudio,
+			URL:     id.ContentURIString(fmt.Sprintf("mxc://%s/%s", resp.ContentURI.Homeserver, resp.ContentURI.FileID)),
+		}
 	}
 
 	_, err = cmd.Matrix.SendEvent(&muksevt.Event{
@@ -296,7 +347,7 @@ func uploadVideo(cmd *Command, filePath string, fileName string, fileSize int, v
 	})
 
 	if err != nil {
-		cmd.Reply("Failed to upload video: %v", err)
+		cmd.Reply("Failed to upload audio/video: %v", err)
 		return
 	}
 }
@@ -329,7 +380,6 @@ func uploadImage(cmd *Command, filePath string, fileName string) {
 	}
 
 	txnID := cmd.Matrix.Client().TxnID()
-	//cmd.Reply("Uploaded image: %v", fmt.Sprintf("mxc://%s/%s", resp.ContentURI.Homeserver, resp.ContentURI.FileID))
 	content := event.MessageEventContent{
 		Body: fileName,
 		Info: &event.FileInfo{
@@ -361,10 +411,81 @@ func uploadImage(cmd *Command, filePath string, fileName string) {
 	}
 }
 
+func uploadFile(cmd *Command, filePath string, fileName string, contentType string) {
+	fileData, err := ioutil.ReadFile(filePath)
+	if err != nil {
+		cmd.Reply("Failed to read file, error: %v", err)
+		return
+	}
+	fileSize := len(fileData)
+
+	resp, err := cmd.Matrix.Client().UploadBytesWithName(fileData, contentType, fileName)
+	if err != nil {
+		cmd.Reply("Failed to upload file: %v", err)
+		return
+	}
+
+	txnID := cmd.Matrix.Client().TxnID()
+	content := event.MessageEventContent{
+		Body: fileName,
+		Info: &event.FileInfo{
+			MimeType: contentType,
+			Size:     int(fileSize),
+		},
+		MsgType: event.MsgFile,
+		URL:     id.ContentURIString(fmt.Sprintf("mxc://%s/%s", resp.ContentURI.Homeserver, resp.ContentURI.FileID)),
+	}
+
+	_, err = cmd.Matrix.SendEvent(&muksevt.Event{
+		Event: &event.Event{
+			Content:   event.Content{Parsed: &content},
+			ID:        id.EventID(txnID),
+			RoomID:    cmd.Room.MxRoom().ID,
+			Timestamp: time.Now().UnixNano() / 1e6,
+			Sender:    cmd.Matrix.Client().UserID,
+			Type:      event.EventMessage,
+			Unsigned:  event.Unsigned{TransactionID: txnID},
+		},
+	})
+
+	if err != nil {
+		cmd.Reply("Failed to upload file: %v", err)
+		return
+	}
+}
+
+func fileGetContentType(filePath string) (string, error) {
+	file, err := os.Open(filePath)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	var buffer [512]byte
+	_, err = file.Read(buffer[0:])
+	if err != nil {
+		return "", err
+	}
+
+	if len(buffer) >= 3 && (buffer[0] == 0xFF && (buffer[1] == 0xFB || buffer[1] == 0xF3 || buffer[1] == 0xF2)) || string(buffer[0:3]) == "ID3" {
+		return "audio/mpeg", nil
+	}
+
+	return http.DetectContentType(buffer[0:]), nil
+}
+
+func isVideoContentType(contentType string) bool {
+	return contentType == "video/avi" || contentType == "video/mp4" || contentType == "video/webm"
+}
+
+func isAudioContentType(contentType string) bool {
+	// audio/mpeg is also mp3
+	return contentType == "audio/basic" || contentType == "audio/aiff" || contentType == "audio/mpeg" || contentType == "audio/midi" || contentType == "audio/wave"
+}
+
 func cmdUpload(cmd *Command) {
 	filePath := cmd.RawArgs
 	fileName := filepath.Base(filePath)
-	fileExt := filepath.Ext(fileName)
 
 	// TODO: This is not safe. Instead open the file and get file size blabla?
 	stat, err := os.Stat(filePath)
@@ -379,17 +500,26 @@ func cmdUpload(cmd *Command) {
 		return
 	}
 
-	var videoFormat string
-	if fileExt == ".webm" {
-		videoFormat = "webm"
-	} else if fileExt == ".mp4" {
-		videoFormat = "mp4"
+	contentType, err := fileGetContentType(filePath)
+	if err != nil {
+		cmd.Reply("Failed to get content type of %s", filePath)
+		return
 	}
 
-	if videoFormat != "" {
-		go uploadVideo(cmd, filePath, fileName, int(fileSize), videoFormat)
+	if isVideoContentType(contentType) || isAudioContentType(contentType) {
+		go uploadAudioVideo(cmd, filePath, fileName, int(fileSize), contentType)
 	} else {
-		go uploadImage(cmd, filePath, fileName)
+		contentType, err := fileGetContentType(filePath)
+		if err != nil {
+			cmd.Reply("Failed to get content type of %s", filePath)
+			return
+		}
+
+		if contentType == "image/jpeg" || contentType == "image/png" || contentType == "image/gif" || contentType == "image/webp" {
+			go uploadImage(cmd, filePath, fileName)
+		} else {
+			go uploadFile(cmd, filePath, fileName, contentType)
+		}
 	}
 }