package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"os/user"
	"path"
	"regexp"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/pbkdf2"
	"golang.org/x/crypto/ssh/terminal"
)

var Categories = map[string]string{
	"Login":            "001",
	"Credit Card":      "002",
	"Secure Note":      "003",
	"Identity":         "004",
	"Password":         "005",
	"Tombstone":        "099",
	"Software License": "100",
	"Bank Account":     "101",
	"Database":         "102",
	"Driver License":   "103",
	"Outdoor License":  "104",
	"Membership":       "105",
	"Passport":         "106",
	"Rewards":          "107",
	"SSN":              "108",
	"Router":           "109",
	"Server":           "110",
	"Email":            "111",
}

var CategoriesRev = map[string]string{
	"001": "Login",
	"002": "Credit Card",
	"003": "Secure Note",
	"004": "Identity",
	"005": "Password",
	"099": "Tombstone",
	"100": "Software License",
	"101": "Bank Account",
	"102": "Database",
	"103": "Driver License",
	"104": "Outdoor License",
	"105": "Membership",
	"106": "Passport",
	"107": "Rewards",
	"108": "SSN",
	"109": "Router",
	"110": "Server",
	"111": "Email",
}

type Key struct {
	Encryption []byte
	Mac        []byte
}

func NewKey(combined []byte) *Key {
	if len(combined) != 64 {
		panic("Invalid key size")
	}

	return &Key{combined[:32], combined[32:64]}
}

func NewKeyPBKDF2(pass []byte, p *OPProfile) *Key {
	return NewKey(pbkdf2.Key(pass, p.Salt, p.Iterations, 64, sha512.New))
}

type ODataOverview struct {
	Title string `json:"title"`
	URL   string `json:"url"`
}

type OPDataSection struct {
	Name   string        `json:"name"`
	Title  string        `json:"title"`
	Fields []OPDataField `json:"fields"`
}

type OPDataField struct {
	ID          string // id
	Name        string // name, n
	Type        string // type, t
	Kind        string // k
	Value       string // value, v
	Designation string // designation
}

type OPBandItem struct {
	Profile         *OPProfile `json:"-"`
	UUID            string     `json:"uuid"`
	Category        string     `json:"category"`
	CategoryName    string     `json:"-"`
	Folder          string     `json:"folder"`
	Trashed         bool       `json:"trashed"`
	Favorite        int        `json:"fave"`
	Data            []byte     `json:"d"`
	HMAC            []byte     `json:"hmac"`
	Key             []byte     `json:"k"`
	Overview        []byte     `json:"o"`
	Created         time.Time  `json:"created"`
	TransactionTime time.Time  `json:"tx"`
	Updated         time.Time  `json:"updated"`
}

func (i *OPBandItem) UnmarshalJSON(data []byte) error {
	var err error
	type LocalItem OPBandItem

	ti := &struct {
		Data            string `json:"d"`
		HMAC            string `json:"hmac"`
		Key             string `json:"k"`
		Overview        string `json:"o"`
		Created         int64  `json:"created"`
		Updated         int64  `json:"updated"`
		TransactionTime int64  `json:"tx"`
		*LocalItem
	}{
		LocalItem: (*LocalItem)(i),
	}

	if err = json.Unmarshal(data, &ti); err != nil {
		return err
	}

	i.Data, err = base64.StdEncoding.DecodeString(ti.Data)
	if err != nil {
		return err
	}

	i.HMAC, err = base64.StdEncoding.DecodeString(ti.HMAC)
	if err != nil {
		return err
	}

	i.Key, err = base64.StdEncoding.DecodeString(ti.Key)
	if err != nil {
		return err
	}

	i.Overview, err = base64.StdEncoding.DecodeString(ti.Overview)
	if err != nil {
		return err
	}

	i.CategoryName = CategoriesRev[i.Category]
	i.Created = time.Unix(ti.Created, 0)
	i.Updated = time.Unix(ti.Updated, 0)
	i.TransactionTime = time.Unix(ti.TransactionTime, 0)

	return nil
}

func (i *OPBandItem) DecryptKey(key *Key) (*Key, error) {
	data := i.Key

	mac := hmac.New(sha256.New, key.Mac)
	mac.Write(data[:len(data)-32])
	expectedMac := mac.Sum(nil)

	if !hmac.Equal(data[len(data)-32:len(data)], expectedMac) {
		return nil, errors.New("HMAC does not match")
	}

	iv := data[:16]
	ct := data[16 : len(data)-32]

	block, err := aes.NewCipher(key.Encryption)
	if err != nil {
		return nil, err
	}

	if len(ct)%aes.BlockSize != 0 {
		return nil, errors.New("Ciphertext is not multiple of block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ct, ct)

	return NewKey(ct), nil
}

func (i *OPBandItem) DecryptData() ([]byte, error) {
	k, err := i.DecryptKey(i.Profile.masterKey)
	if err != nil {
		return nil, err
	}

	d, err := ParseOpdata01(i.Data, k)
	if err != nil {
		return nil, err
	}

	return d, nil
}

func (i *OPBandItem) DecryptOverview() ([]byte, error) {
	o, err := ParseOpdata01(i.Overview, i.Profile.overviewKey)
	if err != nil {
		return nil, err
	}

	return o, nil
}

type OPProfile struct {
	Path          string    `json:"-"`
	UUID          string    `json:"uuid"`
	ProfileName   string    `json:"profileName"`
	Iterations    int       `json:"iterations"`
	Salt          []byte    `json:"salt"`
	MasterKey     []byte    `json:"masterKey"`
	OverviewKey   []byte    `json:"overviewKey"`
	CreatedAt     time.Time `json:"createdAt"`
	UpdatedAt     time.Time `json:"updatedAt"`
	LastUpdatedBy string    `json:"lastUpdatedBy"`

	masterKey   *Key
	overviewKey *Key
	Items       map[string]*OPBandItem `json:"-"`
}

func (p *OPProfile) UnmarshalJSON(data []byte) error {
	var err error
	type LocalProfile OPProfile

	tp := &struct {
		CreatedAt   int64  `json:"createdAt"`
		UpdatedAt   int64  `json:"updatedAt"`
		Salt        string `json:"salt"`
		MasterKey   string `json:"masterKey"`
		OverviewKey string `json:"overviewKey"`
		*LocalProfile
	}{
		LocalProfile: (*LocalProfile)(p),
	}

	if err = json.Unmarshal(data, &tp); err != nil {
		return err
	}

	p.Salt, err = base64.StdEncoding.DecodeString(tp.Salt)
	if err != nil {
		return err
	}

	p.MasterKey, err = base64.StdEncoding.DecodeString(tp.MasterKey)
	if err != nil {
		return err
	}

	p.OverviewKey, err = base64.StdEncoding.DecodeString(tp.OverviewKey)
	if err != nil {
		return err
	}

	p.CreatedAt = time.Unix(tp.CreatedAt, 0)
	p.UpdatedAt = time.Unix(tp.UpdatedAt, 0)

	return nil
}

func (p *OPProfile) Unlock(masterpass []byte) error {
	var err error
	derived := NewKeyPBKDF2(masterpass, p)

	p.masterKey, err = DeriveKey(p.MasterKey, derived)
	if err != nil {
		return err
	}

	p.overviewKey, err = DeriveKey(p.OverviewKey, derived)
	if err != nil {
		return err
	}

	return nil
}

func (p *OPProfile) LoadAllBands() {
	for i := 0; i < 16; i++ {
		_ = p.LoadBand(fmt.Sprintf("%X", i))
	}
}

func (p *OPProfile) LoadBand(band string) error {
	c, err := ioutil.ReadFile(path.Join(p.Path, "default", fmt.Sprintf("band_%s.js", band)))
	if err != nil {
		return err
	}

	// JSON surrounded by "ld({...json...});"
	b := make(map[string]*OPBandItem)
	if err = json.Unmarshal(c[3:len(c)-2], &b); err != nil {
		return err
	}

	for k, v := range b {
		v.Profile = p
		o, err := v.DecryptOverview()
		if err == nil {
			v.Overview = o
		}
		p.Items[k] = v
	}

	return nil
}

func LoadProfile(vaultpath string, masterpass []byte) (*OPProfile, error) {
	p, err := NewProfile(vaultpath)
	if err != nil {
		return nil, err
	}

	err = p.Unlock(masterpass)
	if err != nil {
		return nil, err
	}

	// Wipe master password buffer
	for i, _ := range masterpass {
		masterpass[i] = 0
	}

	p.LoadAllBands()

	return p, nil
}

func NewProfile(vaultpath string) (*OPProfile, error) {
	c, err := ioutil.ReadFile(path.Join(vaultpath, "default", "profile.js"))
	if err != nil {
		return nil, err
	}

	cut := bytes.Index(c, []byte("{"))
	if cut == -1 {
		return nil, errors.New("Profile not a valid JSON document")
	}

	// JSON surrounded by "var profile={...json...};"
	p := &OPProfile{
		Path:  vaultpath,
		Items: make(map[string]*OPBandItem),
	}
	if err = json.Unmarshal(c[cut:len(c)-1], &p); err != nil {
		return nil, err
	}

	return p, nil
}

func ParseOpdata01(data []byte, key *Key) ([]byte, error) {
	// Validate data header
	if !bytes.Equal(data[:8], []byte("opdata01")) {
		return nil, errors.New("opdata01 header mismatch")
	}

	// Validate HMAC before we continue
	mac := hmac.New(sha256.New, key.Mac)
	mac.Write(data[:len(data)-32])
	expectedMac := mac.Sum(nil)

	if !hmac.Equal(data[len(data)-32:len(data)], expectedMac) {
		return nil, errors.New("HMAC does not match")
	}

	iv := data[16:32]
	ct := data[32 : len(data)-32]
	plaintextLen := binary.LittleEndian.Uint64(data[8:16])

	block, err := aes.NewCipher(key.Encryption)
	if err != nil {
		return nil, err
	}

	if len(ct)%aes.BlockSize != 0 {
		return nil, errors.New("Ciphertext is not multiple of block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ct, ct)

	// Copy so we can free the buffer with IV/padding prefix
	out := make([]byte, plaintextLen)
	copy(out, ct[len(ct)-int(plaintextLen):len(ct)])

	return out, nil
}

func DeriveKey(data []byte, dkey *Key) (*Key, error) {
	raw, err := ParseOpdata01(data, dkey)
	if err != nil {
		return nil, err
	}

	h := sha512.New()
	h.Write(raw)

	return NewKey(h.Sum(nil)), nil
}

type SearchResult struct {
	Title    string `json:"title"`
	URL      string `json:"url"`
	BandItem *OPBandItem
}

func SearchOverviews(needle string, haystack map[string]*OPBandItem) (map[string]SearchResult, error) {
	test, err := regexp.Compile(fmt.Sprintf("(?iU).*%s.*", needle))
	if err != nil {
		return nil, err
	}

	m := make(map[string]SearchResult)

	// Simple linear search because the data set size is small
	for k, v := range haystack {
		if v.Trashed {
			continue
		}

		var o SearchResult
		if err = json.Unmarshal(v.Overview, &o); err != nil {
			fmt.Printf("[ERROR] Decoding overview for %s\n", k)
			continue
		}

		if test.MatchString(o.Title) || test.MatchString(o.URL) {
			m[k] = SearchResult{o.Title, o.URL, v}
		}
	}

	return m, nil
}

// TODO: a more sophisticated approach, perhaps
func TruncateURL(u string) string {
	if u == "" {
		return u
	}

	p, err := url.Parse(u)
	if err != nil {
		if len(u) > 20 {
			return fmt.Sprintf("%s...", u[:17])
		} else {
			return u
		}
	} else {
		return fmt.Sprintf("%s://%s", p.Scheme, p.Host)
	}
}

// TODO: testing
func doSearch(p *OPProfile) {
	candidates, err := SearchOverviews(os.Args[1], p.Items)
	if err != nil {
		fmt.Println("Error searching")
		return
	} else {
		fmt.Printf("Found %d candidates\n", len(candidates))

		outs := make([][]string, 0, len(candidates)+1)
		c1len := 0

		for k, v := range candidates {
			c1l := len(v.BandItem.CategoryName)
			if c1l > c1len {
				c1len = c1l
			}
			outs = append(outs, []string{v.BandItem.CategoryName, k, v.Title, TruncateURL(v.URL)})
		}

		for _, i := range outs {
			pad := ""
			if len(i[0]) < (c1len + 5) {
				pad = strings.Join(make([]string, c1len-(len(i[0])-1)), " ")
			}
			fmt.Printf("%s %s %s %s (%s)\n", i[0], pad, i[1], i[2], i[3])
		}

		return
	}
}

// TODO: testing
func printItem(p *OPProfile, idx string) {
	item, ok := p.Items[idx]
	if !ok {
		fmt.Println("Error: No such item")
		return
	}

	itemDetails, err := item.DecryptData()
	if err != nil {
		fmt.Println("Error decoding item data")
		fmt.Println(err)
		return
	}

	fmt.Printf("[\"%s\", %s, %s]", item.CategoryName, item.Overview, itemDetails)
}

func ExpandUser(path string) (string, error) {
	if strings.HasPrefix(datapath, "~/") {
		if u, err := user.Current(); err != nil {
			return nil, err
		} else {
			return fmt.Sprintf("%s%s", u.HomeDir, datapath[1:len(datapath)]), nil
		}
	}
	return path, nil
}

func main() {
	datapath, err := ExpandUser("~/Dropbox/1Password/1Password.opvault")
	if err != nil {
		fmt.Printf(err)
		return
	}

	fmt.Print("Enter Password: ")
	password, err := terminal.ReadPassword(int(syscall.Stdin))
	if err != nil {
		fmt.Println("Unable to read password")
		return
	}
	fmt.Println("")

	p, err := LoadProfile(datapath, password)
	if err != nil {
		fmt.Printf("ERROR: Unable to load keychain: %s", err)
		return
	}

	// TODO: actually handle commands, maybe cobra?
	doSearch(p)
	//printItem(p, "")
}