amuse/db/db.go

package db

import (
	"notabug.org/apiote/amuse/utils"

	"database/sql"
	"errors"
	"fmt"
	"os"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type EmptyError struct {
	message string
}

func (e EmptyError) Error() string {
	return e.message
}

type User struct {
	Username      string
	PasswordHash  string
	Sfa           string
	Avatar        []byte
	AvatarSmall   []byte
	IsAdmin       bool
	RecoveryCodes string
}

type Session struct {
	Id       string
	Username string
	Expiry   time.Time
}

func Migrate() error {
	// todo migrations
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		return err
	}
	defer db.Close()

	_, err = db.Exec(`create table cache(uri text primary key, etag text, date date, response blob, last_hit date)`)
	if err != nil {
		return err
	}
	_, err = db.Exec(`create table users(username text primary key, password text, sfa text, avatar blob, avatar_small blob, is_admin bool, recovery_codes text)`)
	if err != nil {
		return err
	}
	_, err = db.Exec(`create table sessions(id text primary key, username text, expiry datetime, foreign key(username) references users(username))`)
	if err != nil {
		return err
	}

	return nil
}

func MakeAdmin(username string) error {
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		fmt.Fprintf(os.Stderr, "DB open err\n")
		return err
	}
	defer db.Close()
	_, err = db.Exec("update users set is_admin = 1 where username = ?", username)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Update err %v\n", err)
		return err
	}
	rows, err := db.Query(`select is_admin from users where username = ?`, username)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Select err %v\n", err)
		return err
	}
	defer rows.Close()

	if !rows.Next() {
		fmt.Fprintf(os.Stderr, "User %s does not exist\n", username)
		return errors.New("User does not exist")
	}
	var isAdmin bool
	err = rows.Scan(&isAdmin)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Scan err %v\n", err)
		return err
	}
	fmt.Println(isAdmin)
	return nil
}

func GetUser(username string) (*User, error) {
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		fmt.Fprintf(os.Stderr, "DB open err\n")
		return nil, err
	}
	defer db.Close()
	rows, err := db.Query(`select password, sfa, recovery_codes, is_admin, avatar, avatar_small from users where username = ?`, username)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Select err %v\n", err)
		return nil, err
	}
	defer rows.Close()
	if !rows.Next() {
		return nil, EmptyError{message: "User does not exist"}
	}
	user := User{Username: username}
	err = rows.Scan(&user.PasswordHash, &user.Sfa, &user.RecoveryCodes, &user.IsAdmin, &user.Avatar, &user.AvatarSmall)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Scan err %v\n", err)
		return nil, err
	}
	return &user, nil
}

func UpdateRecoveryCodes(username, recoveryCodes string) error {
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		fmt.Fprintf(os.Stderr, "DB open err\n")
		return err
	}
	defer db.Close()

	_, err = db.Exec(`update users set recovery_codes = ? where username = ?`, recoveryCodes, username)
	if err != nil {
		return err
	}

	return nil
}

func CreateSession(username, sessionId string, long bool) error {
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		fmt.Fprintf(os.Stderr, "DB open err\n")
		return err
	}
	defer db.Close()

	var length string
	if long {
		length = "30 days"
	} else {
		length = "1 hour"
	}
	_, err = db.Exec(`insert into sessions values(?, ?, datetime('now', '`+length+`'))`, sessionId, username)
	if err != nil {
		return err
	}

	return nil
}

func GetSession(token string) (*Session, error) {
	db, err := sql.Open("sqlite3", utils.DataHome+"/amuse.db")
	if err != nil {
		fmt.Fprintf(os.Stderr, "DB open err\n")
		return nil, err
	}
	defer db.Close()

	rows, err := db.Query(`select username, expiry from sessions where id = ?`, token)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Select err %v\n", err)
		return nil, err
	}
	defer rows.Close()
	if !rows.Next() {
		return nil, EmptyError{message: "Session does not exist"}
	}
	session := Session{Id: token}
	err = rows.Scan(&session.Username, &session.Expiry)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Scan err %v\n", err)
		return nil, err
	}
	return &session, nil
}