generic-mastodon-authenticator/main.go

package main

// Small generic oauth authentication server to retrieve oauth tokens
// from mastodon instances.

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"github.com/labstack/gommon/log"
	"go.etcd.io/bbolt"
	"gopkg.in/yaml.v2"
)

var (
	upgrader = websocket.Upgrader{}
)

type Instance struct {
	Secret string `json:"client_secret"`
	ID     string `json:"client_id"`
	Host   string
}

type App struct {
	AuthMap *AuthenticationMap
	Logger  *echo.Logger
	DB      *bbolt.DB
	Config  *AppConfig
}

type AppConfig struct {
	AppName   string   `yaml:"app_name"`
	AppHost   string   `yaml:"app_host"`
	AppScheme string   `yaml:"app_scheme"`
	DBPath    string   `yaml:"db_path"`
	AppScopes []string `yaml:"app_scopes"`
	Website   string   `yaml:"website"`
	Interface string   `yaml:"network_interface"`
	Port      string   `yaml:"network_port"`
}

type AuthToken struct {
	RequestID        string
	Token            string `json:"access_token"`
	Error            string `json:"error"`
	ErrorDescription string `json:"error_description"`
}

func (a *App) AuthRequestWebSocket(c echo.Context) error {
	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)

	if err != nil {
		return err
	}

	client := &AuthRequestClient{
		conn: ws,
		app:  a,
	}

	req, err := client.ReceiveRequest()

	if err != nil {
		return err
	}

	(*a.Logger).Info("Queueing AuthRequest")

	a.AuthMap.QueueRequest(req)
	return nil
}

func (a *App) AuthRequestConfirm(c echo.Context) error {
	code := c.QueryParam("code")
	reqID := c.QueryParam("state")

	if code == "" || reqID == "" {
		return c.String(
			http.StatusUnprocessableEntity,
			"Could not process request, missing required parameters",
		)
	}

	authRequest := a.AuthMap.GetRequestByID(reqID)

	if authRequest == nil {
		(*a.Logger).Errorf("Could not find auth request by ID %s", reqID)
		c.String(http.StatusNotFound, "Auth request with given ID not found")
	}

	instance, err := a.GetCredentialsForInstanceHost(authRequest.Instance)

	if err != nil {
		(*a.Logger).Errorf(
			"Was not able to retrieve instance credentials for host %s during confirm action: %s",
			authRequest.Instance,
			err,
		)

		return c.String(
			http.StatusInternalServerError,
			"Something has gone wrong, please note down the current time and contact the server admin",
		)
	}

	token, err := a.GetAuthTokenForAuthorizationCode(code, instance)

	if err != nil {
		(*a.Logger).Errorf(
			"Failed to retreive acces token from instance %s: %s",
			instance.Host,
			err,
		)

		return c.String(
			http.StatusOK,
			"Something went wrong while retrieving the access token. Please contact the server admin",
		)
	}

	err = authRequest.FulFill(token)
	a.AuthMap.Delete(authRequest.ID)

	if err != nil {
		(*a.Logger).Errorf(
			"Unable to fulfill authentication request with ID %s : %s",
			authRequest.ID,
			err,
		)
		return c.String(
			http.StatusInternalServerError,
			"Something went wrong with the connection to your client",
		)
	}

	return c.String(http.StatusOK, "Athentication completed, you can close this window now")
}

func (a *App) GetAuthTokenForAuthorizationCode(code string, instance *Instance) (*AuthToken, error) {
	data := map[string]string{
		"client_id":     instance.ID,
		"client_secret": instance.Secret,
		"code":          code,
		"grant_type":    "authorization_code",
		"redirect_uri":  a.Config.AppScheme + "://" + a.Config.AppHost + "/confirm",
	}

	jsonData, err := json.Marshal(data)

	if err != nil {
		return nil, err
	}

	response, err := http.Post("https://"+instance.Host+"/oauth/token", "application/json", bytes.NewBuffer(jsonData))

	if err != nil {
		return nil, err
	}

	responseData, err := ioutil.ReadAll(response.Body)

	if err != nil {
		return nil, err
	}

	authToken := &AuthToken{}

	err = json.Unmarshal(responseData, authToken)

	if err != nil {
		return nil, err
	}

	if authToken.Error != "" {
		return nil, errors.New(authToken.Error + ": " + authToken.ErrorDescription)
	}

	return authToken, nil
}

func (a *App) GetCredentialsForInstanceHost(host string) (*Instance, error) {
	var instance *Instance

	err := a.DB.View(func(t *bbolt.Tx) error {
		instanceJson := t.Bucket([]byte("instances")).Get([]byte(host))

		if instanceJson == nil {
			instance = nil
			return nil
		}

		instance = &Instance{}

		err := json.Unmarshal(instanceJson, instance)

		return err
	})

	if instance == nil && err == nil {
		(*a.Logger).Infof("No instance data known for host '%s', requesting credentials from host\n", host)

		instance, err = a.GetCredentialsFromInstanceHost(host)

		if err != nil {
			return nil, err
		}

		err = a.DB.Update(func(t *bbolt.Tx) error {
			instanceJson, err := json.Marshal(instance)

			if err != nil {
				return err
			}

			err = t.Bucket([]byte("instances")).Put([]byte(host), instanceJson)

			return err
		})
	}

	return instance, err
}

func (a *App) AuthRequestRedirect(c echo.Context) error {
	requestId := c.Param("request_id")

	authRequest := a.AuthMap.GetRequestByID(requestId)

	if authRequest == nil {
		return c.String(http.StatusNotFound, "No authentication request found by this ID")
	}

	instance, err := a.GetCredentialsForInstanceHost(authRequest.Instance)

	if err != nil {
		(*a.Logger).Errorf(
			"Failed to get credentials for instance host %s: %s",
			authRequest.Instance,
			err,
		)
	}

	if instance == nil {
		return c.String(
			http.StatusInternalServerError,
			"Something has gone wrong, please note down the current time and contact the server admin",
		)
	}

	values := &url.Values{
		"client_id":     []string{instance.ID},
		"redirect_uri":  []string{a.Config.AppScheme + "://" + a.Config.AppHost + "/confirm"},
		"scope":         []string{strings.Join(a.Config.AppScopes, " ")},
		"response_type": []string{"code"},
		"state":         []string{authRequest.ID},
	}

	return c.Redirect(
		http.StatusTemporaryRedirect,
		"https://"+authRequest.Instance+"/oauth/authorize?"+values.Encode(),
	)
}

func (c *AppConfig) FromFile(filePath string) error {
	yamlfile, err := ioutil.ReadFile(filePath)

	if err != nil {
		return err
	}

	err = yaml.Unmarshal(yamlfile, c)

	if err != nil {
		return err
	}

	if c.AppName == "" || c.AppHost == "" || c.DBPath == "" || c.Website == "" {
		return errors.New("app_name, app_host, db_path and website are required config parameters")
	}

	if c.AppScheme == "" {
		c.AppScheme = "https"
	}

	if c.Port == "" {
		c.Port = "9569"
	}

	if c.Interface == "" {
		c.Interface = "127.0.0.1"
	}

	return nil
}

func (a *App) InitializeDB() error {
	(*a.Logger).Info("Initializing DB from file: ", a.Config.DBPath)

	db, err := bbolt.Open(a.Config.DBPath, 0600, nil)

	if err != nil {
		return err
	}

	err = db.Update(func(t *bbolt.Tx) error {
		_, err := t.CreateBucketIfNotExists([]byte("instances"))

		return err
	})

	a.DB = db
	return nil
}

func (a *App) GetCredentialsFromInstanceHost(host string) (*Instance, error) {
	postData := &map[string]string{
		"client_name":   a.Config.AppName,
		"redirect_uris": a.Config.AppScheme + "://" + a.Config.AppHost + "/confirm",
		"scopes":        strings.Join(a.Config.AppScopes, " "),
		"website":       a.Config.Website,
	}

	data, err := json.Marshal(postData)

	if err != nil {
		return nil, err
	}

	response, err := http.Post("https://"+host+"/api/v1/apps", "application/json", bytes.NewBuffer(data))

	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(response.Body)

	if err != nil {
		return nil, err
	}

	instance := &Instance{
		Host: host,
	}

	err = json.Unmarshal(body, instance)

	return instance, err
}

func main() {
	arguments := os.Args[1:]

	if len(arguments) != 1 {
		fmt.Fprintln(os.Stderr, "Expected 1 command line argument, but got :  ", len(arguments))
		fmt.Fprintln(os.Stderr, "Usage: generic-mastodon-authenticator CONFIG_PATH")
		os.Exit(1)
	}

	config := &AppConfig{}

	err := config.FromFile(arguments[0])

	if err != nil {
		log.Fatal(err)
	}

	e := echo.New()
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	e.Static("/", "../public")

	e.Logger.SetLevel(log.INFO)

	a := &App{
		Logger: &e.Logger,
		Config: config,
	}

	m := &AuthenticationMap{
		app: a,
	}

	m.Init()

	a.AuthMap = m

	err = a.InitializeDB()

	if err != nil {
		log.Fatal(err)
	}

	e.GET("/ws", a.AuthRequestWebSocket)
	e.GET("/auth/:request_id", a.AuthRequestRedirect)
	e.GET("/confirm", a.AuthRequestConfirm)

	e.Logger.Fatal(e.Start(a.Config.Interface + ":" + a.Config.Port))
}