drivel/media.go

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

const (
	MAX_BYTES       = 50 * 1024 * 1024
	CHUNK_SIZE      = 1024 * 1024
	UPLOAD_ENDPOINT = "https://upload.twitter.com/1.1/media/upload.json"
)

func InitRequest(mediaType string, totalBytes int) url.Values {
	r := map[string][]string{
		"command":     {"INIT"},
		"media_type":  {mediaType},
		"total_bytes": {strconv.Itoa(totalBytes)},
	}
	if mediaType == "video/mp4" {
		r["media_category"] = []string{"tweet_video"}
	}
	return r
}

type InitResponse struct {
	Errors          []TwitterError
	Media_id_string string
}

func (ir InitResponse) Error() string {
	if len(ir.Errors) == 0 {
		return ""
	} else {
		s, _ := json.Marshal(ir)
		return "Response error " + string(s)
	}
}

func AppendRequest(mediaID ObjectID, mediaData string, segmentIndex int) url.Values {
	return map[string][]string{
		"command":       {"APPEND"},
		"media_id":      {string(mediaID)},
		"media_data":    {mediaData},
		"segment_index": {strconv.Itoa(segmentIndex)},
	}
}

func FinalizeRequest(mediaID ObjectID) url.Values {
	return map[string][]string{
		"command":  {"FINALIZE"},
		"media_id": {string(mediaID)},
	}
}

type FinalizeResponse struct {
	Error           string
	Media_id_string string
	Processing_info ProcessingInfo
}

type ProcessingInfo struct {
	State            string
	Check_after_secs int64
	Progress_percent int64
	Error            TwitterError
}

func PollStatusParameters(mediaID ObjectID) string {
	return "?command=STATUS&media_id=" + string(mediaID)
}

type PollStatusResponse struct {
	Processing_info ProcessingInfo
}

var mimetype = map[string]string{
	".mp4":  "video/mp4",
	".jpg":  "image/jpeg",
	".jpeg": "image/jpeg",
	".png":  "image/png",
	".gif":  "image/gif",
}

func getMimetype(file string) string {
	ext := filepath.Ext(file)
	if v, ok := mimetype[ext]; ok {
		return v
	} else {
		return "application/octet-stream"
	}
}

func initFileUpload(file string, mediaData []byte) ObjectID {
	log := func(err error) { optLogFatal("initFileUpload "+file, err) }
	initRequest := InitRequest(getMimetype(file), len(mediaData))
	body := post(UPLOAD_ENDPOINT, initRequest)
	var initResponse InitResponse
	err := json.Unmarshal(body, &initResponse)
	log(err)
	log(initResponse)
	return ObjectID(initResponse.Media_id_string)
}

func appendFileChunks(file string, media string, mediaId ObjectID) {
	log := func(err error) { optLogFatal("appendFileChunks", err) }
	info := func(v ...interface{}) {
		if len(media) > CHUNK_SIZE {
			fmt.Fprintln(os.Stderr, v...)
		}
	}
	info("chunk upload", file)
	info("total", len(media))
	for i := 0; i*CHUNK_SIZE < len(media); i = i + 1 {
		start := i * CHUNK_SIZE
		end := (i + 1) * CHUNK_SIZE
		if end > len(media) {
			end = len(media)
		}
		info("segment", i, "start", start, "end", end)
		appended := false
		var body []byte
		for try := 0; try < 3 && !appended; try++ {
			appRequest := AppendRequest(mediaId, media[start:end], i)
			body = post(UPLOAD_ENDPOINT, appRequest)
			if string(body) == "" {
				appended = true
			}
		}
		if !appended {
			log(errors.New(string(body)))
		}
	}
}

func finalizeFileUpload(file string, mediaId ObjectID) int64 {
	log := func(err error) { optLogFatal("finalizeFileUpload", err) }
	body := post(UPLOAD_ENDPOINT, FinalizeRequest(mediaId))
	var finalizeResponse FinalizeResponse
	err := json.Unmarshal(body, &finalizeResponse)
	log(err)
	log(errors.New(finalizeResponse.Error))
	if id := ObjectID(finalizeResponse.Media_id_string); id != "" {
		fmt.Fprintln(os.Stderr, "==> Uploaded "+file+" with id "+string(id))
		procInfo := finalizeResponse.Processing_info
		return procInfo.Check_after_secs
	} else {
		log(errors.New("Could not finalize " + string(mediaId)))
		return 0
	}
}

func wait(seconds int64) {
	fmt.Fprintln(os.Stderr, "Waiting", seconds, "seconds")
	time.Sleep(time.Duration(seconds) * time.Second)
}

func pollStatus(mediaId ObjectID) {
	log := func(err error) { optLogFatal("pollStatus "+string(mediaId), err) }
	succeeded := false
	var error TwitterError
	for try := 0; try < 6; try = try + 1 {
		body := get(UPLOAD_ENDPOINT + PollStatusParameters(mediaId))
		var response PollStatusResponse
		err := json.Unmarshal(body, &response)
		log(err)
		procInfo := response.Processing_info
		state := procInfo.State
		error = procInfo.Error
		if state == "succeeded" {
			succeeded = true
			break
		} else if state == "failed" {
			break
		} else {
			fmt.Fprintln(os.Stderr, "Processing progress: ", procInfo.Progress_percent, "%")
			seconds := procInfo.Check_after_secs
			if seconds > 10 {
				seconds = 10
			}
			wait(seconds)
		}
	}
	if !succeeded {
		log(errors.New("File upload failed " + error.Message))
	}
}

func uploadFile(file string) ObjectID {
	log := func(err error) { optLogFatal("uploadFile "+file, err) }
	tmpMedia, err := ioutil.ReadFile(file)
	log(err)
	media := base64.RawURLEncoding.EncodeToString(tmpMedia)
	mediaId := initFileUpload(file, tmpMedia)
	appendFileChunks(file, media, mediaId)
	seconds := finalizeFileUpload(file, mediaId)
	if seconds > 0 {
		wait(seconds)
		pollStatus(mediaId)
	}
	return mediaId
}

func uploadAll(files []string) []ObjectID {
	ids := []ObjectID{}
	for _, f := range files {
		if f != "" {
			id := uploadFile(f)
			ids = append(ids, id)
		}
	}
	return ids
}

type mediaKind int

const (
	UNKNOWN mediaKind = iota
	PIC
	GIF
	VIDEO
)

func kind(path string) mediaKind {
	ext := filepath.Ext(path)
	switch ext {
	case ".jpg":
		fallthrough
	case ".jpeg":
		fallthrough
	case ".png":
		return PIC
	case ".gif":
		return GIF
	case ".mp4":
		return VIDEO
	}
	return UNKNOWN
}

func exitIfInvalid(path string) {
	log := func(err error) { optLogFatal("exitIfInvalid", err) }
	// check existence AND readability
	f, err := os.Open(path)
	log(err)
	defer f.Close()
	tmp, err := ioutil.ReadAll(f)
	log(err)
	if len(tmp) > MAX_BYTES {
		log(errors.New("File too big: " + path + " is bigger than maximum of " + strconv.Itoa(MAX_BYTES) + " Bytes"))
	}
}