drivel/drivel.go

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"git.gutmet.org/goutil.git"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	MAX_BYTES         = 50 * 1024 * 1024
	CHUNK_SIZE        = 1024 * 1024
	CHARACTER_LIMIT   = 280
	UPLOAD_ENDPOINT   = "https://upload.twitter.com/1.1/media/upload.json"
	STATUS_ENDPOINT   = "https://api.twitter.com/1.1/statuses/update.json"
	MENTIONS_ENDPOINT = "https://api.twitter.com/1.1/statuses/mentions_timeline.json?tweet_mode=extended&count=200"
	HOME_ENDPOINT     = "https://api.twitter.com/1.1/statuses/home_timeline.json?tweet_mode=extended&count=200"
	LOOKUP_ENDPOINT   = "https://api.twitter.com/1.1/statuses/lookup.json?tweet_mode=extended"
	RETWEET_ENDPOINT  = "https://api.twitter.com/1.1/statuses/retweet/"
	LIKE_ENDPOINT     = "https://api.twitter.com/1.1/favorites/create.json"
)

func optLogFatal(decorum string, err error) {
	if err != nil && err.Error() != "" {
		fmt.Fprintln(os.Stderr, "drivel: "+decorum+": "+err.Error())
		os.Exit(-1)
	}
}

type ObjectID string

type TwitterError struct {
	Code    int64
	Message string
	Label   string
}

type Response struct {
	Errors []TwitterError
}

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

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 {
	Response
	Media_id_string string
}

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
}

func UpdateStatusRequest(status string, mediaIDs []ObjectID, previousStatusID ObjectID) url.Values {
	r := map[string][]string{"status": {status}}
	if len(mediaIDs) > 0 {
		ids := []string{}
		for _, id := range mediaIDs {
			ids = append(ids, string(id))
		}
		r["media_ids"] = []string{strings.Join(ids, ",")}
	}
	if len(previousStatusID) > 0 {
		r["in_reply_to_status_id"] = []string{string(previousStatusID)}
		r["auto_populate_reply_metadata"] = []string{"true"}
	}
	return r
}

type UpdateStatusResponse struct {
	Response
	Id_str string
}

func LookupParameters(ids []string) string {
	return "&id=" + strings.Join(ids, ",")
}

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 get(client *http.Client, url string) []byte {
	return _send(client, url, nil, false)
}

func send(client *http.Client, url string, vals url.Values) []byte {
	return _send(client, url, vals, true)
}

func _send(client *http.Client, url string, vals url.Values, usePost bool) []byte {
	log := func(err error) {
		v, _ := json.Marshal(vals)
		optLogFatal("get/post "+url+" "+string(v), err)
	}
	var resp *http.Response
	var err error
	if usePost {
		resp, err = client.PostForm(url, vals)
	} else {
		resp, err = client.Get(url)
	}
	log(err)
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		fmt.Fprintln(os.Stderr, "response:", resp, "\n")
		fmt.Fprintln(os.Stderr, "body:", string(body), "\n")
		log(errors.New("HTTP status " + fmt.Sprint(resp.StatusCode)))
	}
	log(err)
	return body
}

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

func appendFileChunks(client *http.Client, file string, media string, mediaId ObjectID) {
	log := func(err error) { optLogFatal("appendFileChunks", err) }
	info := func(v ...interface{}) {
		if len(media) > CHUNK_SIZE {
			fmt.Println(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 = send(client, UPLOAD_ENDPOINT, appRequest)
			if string(body) == "" {
				appended = true
			}
		}
		if !appended {
			log(errors.New(string(body)))
		}
	}
}

func finalizeFileUpload(client *http.Client, file string, mediaId ObjectID) int64 {
	log := func(err error) { optLogFatal("finalizeFileUpload", err) }
	body := send(client, 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.Println("==> 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.Println("Waiting", seconds, "seconds")
	time.Sleep(time.Duration(seconds) * time.Second)
}

func pollStatus(client *http.Client, 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(client, 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.Println("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(client *http.Client, 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(client, file, tmpMedia)
	appendFileChunks(client, file, media, mediaId)
	seconds := finalizeFileUpload(client, file, mediaId)
	if seconds > 0 {
		wait(seconds)
		pollStatus(client, mediaId)
	}
	return mediaId
}

func uploadAll(client *http.Client, files []string) []ObjectID {
	ids := []ObjectID{}
	for _, f := range files {
		if f != "" {
			id := uploadFile(client, 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 splitStatus(status string) []string {
	split := []string{}
	words := strings.Split(status, " ")
	s := ""
	for _, word := range words {
		if s == "" && len(word) <= CHARACTER_LIMIT {
			s = word
		} else if len(s)+1+len(word) <= CHARACTER_LIMIT {
			s = s + " " + word
		} else {
			split = append(split, s)
			bound := goutil.IntMin(len(word), CHARACTER_LIMIT)
			s = string([]rune(word)[:bound])
		}
	}
	if s != "" {
		split = append(split, s)
	}
	return split
}

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"))
	}
}

func splitArguments(args []string) data {
	if len(args) < 1 {
		fmt.Fprintln(os.Stderr, "Usage: drivel status STATUS [FILE1, FILE2, ...]")
		os.Exit(-1)
	}
	d := data{}
	d.status = splitStatus(args[0])
	for _, arg := range args[1:] {
		exitIfInvalid(arg)
		switch kind(arg) {
		case PIC:
			d.pics = append(d.pics, arg)
		case GIF:
			d.gifs = append(d.gifs, arg)
		case VIDEO:
			d.videos = append(d.videos, arg)
		default:
			optLogFatal("splitArguments", errors.New("Unsupported file: "+arg))
		}
	}
	return d
}

func tweet(client *http.Client, status string, mediaIDs []ObjectID, previousID ObjectID) ObjectID {
	log := func(err error) { optLogFatal("tweet "+status, err) }
	request := UpdateStatusRequest(status, mediaIDs, previousID)
	body := send(client, STATUS_ENDPOINT, request)
	var sr UpdateStatusResponse
	err := json.Unmarshal(body, &sr)
	log(err)
	if len(sr.Errors) > 0 {
		log(sr)
	}
	fmt.Println("==> Updated status to '" + status + "' with id " + sr.Id_str)
	return ObjectID(sr.Id_str)
}

type data struct {
	status []string
	pics   []string
	gifs   []string
	videos []string
}

func (d *data) getStatus(i int) string {
	return goutil.StrSliceAt(d.status, i)
}

func (d *data) getGif(i int) string {
	return goutil.StrSliceAt(d.gifs, i)
}

func (d *data) getVideo(i int) string {
	return goutil.StrSliceAt(d.videos, i)
}

func (d *data) uploadPics(client *http.Client, from, to int) []ObjectID {
	pics := goutil.StrSlice(d.pics, from, to)
	return uploadAll(client, pics)
}

func (d *data) uploadGif(client *http.Client, i int) []ObjectID {
	gif := d.getGif(i)
	return uploadAll(client, []string{gif})
}

func (d *data) uploadVideo(client *http.Client, i int) []ObjectID {
	vid := d.getVideo(i)
	return uploadAll(client, []string{vid})
}

func (d *data) push(client *http.Client, previous ObjectID) {
	if d == nil {
		return
	}
	empty := false
	i, g, v := 0, 0, 0
	for !empty {
		empty = true
		status := d.getStatus(i)
		mediaIDs := []ObjectID{}
		if status != "" {
			empty = false
		}
		from := i * 4
		to := (i + 1) * 4
		mediaIDs = d.uploadPics(client, from, to)
		if len(mediaIDs) == 0 {
			mediaIDs = d.uploadGif(client, g)
			g++
		}
		if len(mediaIDs) == 0 {
			mediaIDs = d.uploadVideo(client, v)
			v++
		}
		if len(mediaIDs) > 0 {
			empty = false
		}
		if !empty {
			previous = tweet(client, status, mediaIDs, previous)
			i++
		}
	}
}

func updateStatus(args []string, previous ObjectID, embedTweet ObjectID) {
	d := splitArguments(args)
	httpClient := getClient()
	if embedTweet != "" {
		tweets := lookup(httpClient, []string{string(embedTweet)})
		if len(tweets) == 1 {
			d.status[0] += " " + tweets[0].URL()
		}
	}
	d.push(httpClient, previous)
}

func status(args []string) error {
	updateStatus(args, "", "")
	return nil
}

func reply(args []string) error {
	if len(args) < 2 {
		fmt.Fprintln(os.Stderr, "Usage: drivel reply TWEET_ID MESSAGE [FILE1, FILE2, ...]")
		os.Exit(-1)
	}
	updateStatus(args[1:], ObjectID(args[0]), "")
	return nil
}

func quote(args []string) error {
	if len(args) < 2 {
		fmt.Println(os.Stderr, "Usage: drivel quote TWEET_ID MESSAGE [FILE1, FILE2, ...]")
		os.Exit(-1)
	}
	updateStatus(args[1:], "", ObjectID(args[0]))
	return nil
}

type Status struct {
	Full_text string
	Id_str    string
	User      StatusUser
}

func (m Status) String() string {
	return m.User.Name + " " + "(" + m.Id_str + ")" + ":\n" + m.Full_text
}

func (m Status) URL() string {
	return "https://twitter.com/" + m.User.Screen_name + "/status/" + m.Id_str
}

type StatusUser struct {
	Name        string
	Screen_name string
}

func lookup(client *http.Client, ids []string) []Status {
	log := func(err error) { optLogFatal("lookup "+strings.Join(ids, ","), err) }
	body := get(client, LOOKUP_ENDPOINT+LookupParameters(ids))
	var tweets []Status
	err := json.Unmarshal(body, &tweets)
	log(err)
	return tweets
}

func timeline(endpoint string) {
	log := func(err error) { optLogFatal("timeline", err) }
	client := getClient()
	body := get(client, endpoint)
	var tweets []Status
	err := json.Unmarshal(body, &tweets)
	log(err)
	for _, tweet := range tweets {
		fmt.Println(tweet)
		fmt.Println("---------")
	}
}

func mentions(args []string) error {
	timeline(MENTIONS_ENDPOINT)
	return nil
}

func home(args []string) error {
	timeline(HOME_ENDPOINT)
	return nil
}

func RetweetParameters(id string) string {
	return id + ".json"
}

func retweet(args []string) error {
	log := func(err error) { optLogFatal("retweet", err) }
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "USAGE: drivel retweet TWEET_ID")
		os.Exit(-1)
	}
	client := getClient()
	id := args[0]
	tweets := lookup(client, []string{id})
	if len(tweets) != 1 {
		log(errors.New("Could not find tweet " + id))
	}
	body := send(client, RETWEET_ENDPOINT+RetweetParameters(id), nil)
	var retweet Status
	err := json.Unmarshal(body, &retweet)
	log(err)
	fmt.Println("Retweeted", tweets[0])
	return nil
}

func LikeRequest(id string) url.Values {
	return map[string][]string{
		"id":         {id},
		"tweet_mode": {"extended"},
	}
}

func like(args []string) error {
	log := func(err error) { optLogFatal("like", err) }
	if len(args) != 1 {
		fmt.Fprintln(os.Stderr, "USAGE: drivel like TWEET_ID")
		os.Exit(-1)
	}
	client := getClient()
	body := send(client, LIKE_ENDPOINT, LikeRequest(args[0]))
	var tweet Status
	err := json.Unmarshal(body, &tweet)
	log(err)
	fmt.Println("Liked", tweet)
	return nil
}

func main() {
	commands := []goutil.Command{
		goutil.NewCommand("status", status, "post a status with message and/or media"),
		goutil.NewCommand("home", home, "get your home timeline"),
		goutil.NewCommand("mentions", mentions, "get your mention timeline"),
		goutil.NewCommand("reply", reply, "reply to a tweet with a specific ID"),
		goutil.NewCommand("quote", quote, "quote retweet a tweet with a specific ID"),
		goutil.NewCommand("retweet", retweet, "retweet a tweet with a specific ID"),
		goutil.NewCommand("like", like, "like a tweet with a specific ID"),
	}
	err := goutil.Execute(commands)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}