go-registry/reg.go

package registry

import (
	"errors"
	"net"
	"time"

	tea "github.com/fumiama/gofastTEA"
	"github.com/wdvxdr1123/ZeroBot/utils/helper"
)

type Regedit struct {
	conn net.Conn
	addr string
	tp   tea.TEA
	ts   *tea.TEA
	seq  byte
	buf  [255]byte
}

func NewRegedit(addr, pwd, sps string) *Regedit {
	var tp, ts [16]byte
	if len(pwd) > 15 {
		pwd = pwd[:15]
	}
	if len(sps) > 15 {
		sps = sps[:15]
	}
	copy(tp[:], helper.StringToBytes(pwd))
	copy(ts[:], helper.StringToBytes(sps))
	s := tea.NewTeaCipherLittleEndian(ts[:])
	return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:]), ts: &s}
}

func NewRegReader(addr, pwd string) *Regedit {
	var tp [16]byte
	if len(pwd) > 15 {
		pwd = pwd[:15]
	}
	copy(tp[:], helper.StringToBytes(pwd))
	return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:])}
}

func (r *Regedit) Connect() (err error) {
	r.conn, err = net.Dial("tcp", r.addr)
	return
}

func (r *Regedit) ConnectIn(wait time.Duration) (err error) {
	r.conn, err = net.DialTimeout("tcp", r.addr, wait)
	return
}

func (r *Regedit) Close() (err error) {
	p := NewCmdPacket(CMDEND, []byte("fill"), &r.tp)
	r.conn.Write(p.Encrypt(r.seq))
	r.seq = 0
	return r.conn.Close()
}

func (r *Regedit) Get(key string) (string, error) {
	if len(key) > 127 {
		return "", errors.New("get key too long")
	}
	p := NewCmdPacket(CMDGET, helper.StringToBytes(key), &r.tp)
	r.conn.Write(p.Encrypt(r.seq))
	r.seq++
	ack, err := r.ack()
	if err != nil {
		return "", err
	}
	ackbytes := ack.Decrypt(r.seq)
	if ackbytes == nil {
		return "", errors.New("decrypt ack error")
	}
	a := helper.BytesToString(ackbytes)
	r.seq++
	if a == "erro" {
		return "", errors.New("server ack error")
	}
	if a == "null" {
		a = ""
	}
	return a, nil
}

func (r *Regedit) Set(key, value string) error {
	if r.ts == nil {
		return errors.New("permission denied")
	}
	if len(key) > 127 {
		return errors.New("set key too long")
	}
	if len(value) > 127 {
		return errors.New("set val too long")
	}
	p := NewCmdPacket(CMDSET, helper.StringToBytes(key), r.ts)
	r.conn.Write(p.Encrypt(r.seq))
	r.seq++
	ack, err := r.ack()
	if err != nil {
		return err
	}
	ackbytes := ack.Decrypt(r.seq)
	if ackbytes == nil {
		return errors.New("decrypt ack error")
	}
	a := helper.BytesToString(ackbytes)
	r.seq++
	if a == "erro" {
		return errors.New("server ack error")
	}
	if a != "data" {
		return errors.New("unknown ack error")
	}
	p = NewCmdPacket(CMDDAT, helper.StringToBytes(value), r.ts)
	r.conn.Write(p.Encrypt(r.seq))
	r.seq++
	ack, err = r.ack()
	if err != nil {
		return err
	}
	ackbytes = ack.Decrypt(r.seq)
	if ackbytes == nil {
		return errors.New("decrypt ack error")
	}
	a = helper.BytesToString(ackbytes)
	r.seq++
	if a == "erro" {
		return errors.New("server ack error")
	}
	if a != "succ" {
		return errors.New("unknown ack error")
	}
	return nil
}

func (r *Regedit) Del(key string) error {
	if r.ts == nil {
		return errors.New("permission denied")
	}
	if len(key) > 127 {
		return errors.New("get key too long")
	}
	p := NewCmdPacket(CMDDEL, helper.StringToBytes(key), r.ts)
	r.conn.Write(p.Encrypt(r.seq))
	r.seq++
	ack, err := r.ack()
	if err != nil {
		return err
	}
	ackbytes := ack.Decrypt(r.seq)
	if ackbytes == nil {
		return errors.New("decrypt ack error")
	}
	a := helper.BytesToString(ackbytes)
	r.seq++
	if a == "erro" {
		return errors.New("server ack error")
	}
	if a == "null" {
		return errors.New("no such key")
	}
	if a != "succ" {
		return errors.New("unknown ack error")
	}
	return nil
}

func (r *Regedit) ack() (*CmdPacket, error) {
	n, err := r.conn.Read(r.buf[:])
	if err != nil {
		return nil, err
	}
	for n < 1+1+16 {
		m, err := r.conn.Read(r.buf[n:])
		if err != nil {
			return nil, err
		}
		n += m
	}
	for n < 1+1+16+int(r.buf[1]) {
		m, err := r.conn.Read(r.buf[n:])
		if err != nil {
			return nil, err
		}
		n += m
	}
	return ParseCmdPacket(r.buf[:], &r.tp), nil
}