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

const (
	MAX_BYTES         = 5 * 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=100"
	HOME_ENDPOINT     = "https://api.twitter.com/1.1/statuses/home_timeline.json?tweet_mode=extended&count=100"
)

func optLogFatal(decorum string, err error) {
	if err != nil {
		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 {
	return map[string][]string{
		"command":     {"INIT"},
		"media_type":  {mediaType},
		"total_bytes": {strconv.Itoa(totalBytes)},
	}
}

type InitResponse struct {
	Response
	Media_id_string string
}

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

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

type FinalizeResponse struct {
	Error           string
	Media_id_string string
}

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
}

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("send "+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 uploadFile(client *http.Client, file string) ObjectID {
	log := func(err error) { optLogFatal("uploadFile "+file, err) }
	media, err := ioutil.ReadFile(file)
	log(err)
	initRequest := InitRequest(getMimetype(file), len(media))
	body := send(client, UPLOAD_ENDPOINT, initRequest)
	var initResponse InitResponse
	err = json.Unmarshal(body, &initResponse)
	log(err)
	if len(initResponse.Errors) == 0 {
		mediaId := initResponse.Media_id_string
		appRequest := AppendRequest(mediaId, base64.RawURLEncoding.EncodeToString(media), 0)
		body := send(client, UPLOAD_ENDPOINT, appRequest)
		if string(body) == "" {
			body := send(client, UPLOAD_ENDPOINT, FinalizeRequest(mediaId))
			var finalizeResponse FinalizeResponse
			json.Unmarshal(body, &finalizeResponse)
			if id := ObjectID(finalizeResponse.Media_id_string); id != "" {
				fmt.Println("==> Uploaded " + file + " with id " + string(id))
				return id
			}
		}
	}
	log(errors.New("Could not upload file " + file))
	return ""
}

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) {
	d := splitArguments(args)
	httpClient := getClient()
	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
}

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
}

type StatusUser struct {
	Name string
}

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 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"),
	}
	err := goutil.Execute(commands)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}