sparrow/services/registry/registry.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"github.com/gogf/gf/util/guid"
	"sparrow/pkg/generator"
	"sparrow/pkg/models"
	"sparrow/pkg/rpcs"
)

const (
	flagAESKey = "aeskey"
)

var confAESKey = flag.String(flagAESKey, "", "use your own aes encryting key.")

var errDbActFailt = errors.New("数据库操作失败")

type Registry struct {
	keygen *generator.KeyGenerator
}

// NewRegistry create new registry
func NewRegistry() (*Registry, error) {
	gen, err := generator.NewKeyGenerator(*confAESKey)
	if err != nil {
		return nil, err
	}
	return &Registry{
		keygen: gen,
	}, nil
}

func setVendor(target *models.Vendor, src *models.Vendor) {
	target.ID = src.ID
	target.VendorName = src.VendorName
	target.VendorDescription = src.VendorDescription
	target.VendorKey = src.VendorKey
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

func setProduct(target *models.Product, src *models.Product) {
	target.ID = src.ID
	target.ProductName = src.ProductName
	target.ProductDescription = src.ProductDescription
	target.ProductKey = src.ProductKey
	target.ProductConfig = src.ProductConfig
	target.VendorID = src.VendorID
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
	target.RecordId = src.RecordId
}

func setApplication(target *models.Application, src *models.Application) {
	target.ID = src.ID
	target.AppName = src.AppName
	target.AppDescription = src.AppDescription
	target.AppKey = src.AppKey
	target.ReportUrl = src.ReportUrl
	target.AppToken = src.AppToken
	target.AppDomain = src.AppDomain
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
}

func setDevice(target *models.Device, src *models.Device) {
	target.ID = src.ID
	target.RecordId = src.RecordId
	target.ProductID = src.ProductID
	target.DeviceIdentifier = src.DeviceIdentifier
	target.DeviceSecret = src.DeviceSecret
	target.DeviceKey = src.DeviceKey
	target.DeviceName = src.DeviceName
	target.DeviceDescription = src.DeviceDescription
	target.DeviceVersion = src.DeviceVersion
	target.CreatedAt = src.CreatedAt
	target.UpdatedAt = src.UpdatedAt
	target.VendorID = src.VendorID
}

//func setSubDevice(target *models.SubDevice, src *models.SubDevice) {
//	target.ID = src.ID
//	target.RecordId = src.RecordId
//	target.SubDeviceID = src.SubDeviceID
//	target.DeviceCode = src.DeviceCode
//	target.Status = src.Status
//	target.Name = src.Name
//}

// SaveVendor will create a vendor if the ID field is not initialized
// if ID field is initialized, it will update the conresponding vendor.
func (r *Registry) SaveVendor(vendor *models.Vendor, reply *models.Vendor) error {
	vendor.RecordId = guid.S()
	db, err := getDB()
	if err != nil {
		return err
	}

	if vendor.ID == 0 {
		// if ID field is not initialized, will generate key first
		err = db.Save(vendor).Error
		if err != nil {
			return err
		}

		key, err := r.keygen.GenRandomKey(vendor.RecordId)
		if err != nil {
			return err
		}

		vendor.VendorKey = key
	}

	err = db.Save(vendor).Error
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Vendor:%v", vendor.ID)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}
	// write cache
	cache.Set(cacheKey, vendor)
	setVendor(reply, vendor)

	return nil
}

// SaveApplication will create a application if the ID field is not initialized
// if ID field is initialized, it will update the conresponding application.
func (r *Registry) SaveApplication(app *models.Application, reply *models.Application) error {
	app.RecordId = guid.S()
	db, err := getDB()
	if err != nil {
		return err
	}

	if app.ID == 0 {
		err = db.Save(app).Error
		if err != nil {
			return err
		}

		key, err := r.keygen.GenRandomKey(app.RecordId)
		if err != nil {
			return err
		}

		app.AppKey = key
	}

	err = db.Save(app).Error
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", app.ID)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}

	setApplication(reply, app)

	return nil
}

// ValidateApplication try to validate the given app key.
// if success, it will reply the corresponding application
func (r *Registry) ValidateApplication(key string, reply *models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	id, err := r.keygen.DecodeIDFromRandomKey(key)
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", id)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		app := cacheValue.(*models.Application)
		setApplication(reply, app)
	} else {
		err = db.First(reply, "record_id = ?", id).Error
		if err != nil {
			return err
		}
		var storage models.Application
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	if reply.AppKey != key {
		return errors.New("app key not match")
	}

	return nil
}

// FindVendor will find product by specified ID
func (r *Registry) FindVendor(recordId string, reply *models.Vendor) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Vendor:%v", recordId)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		vendor := cacheValue.(*models.Vendor)
		setVendor(reply, vendor)
	} else {
		err = db.First(reply, "record_id = ?", recordId).Error
		if err != nil {
			return err
		}
		var storage models.Vendor
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	return nil
}

// GetVendors will get all vendors in the platform.
func (r *Registry) GetVendors(noarg int, reply *[]*models.Vendor) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	return db.Find(reply).Error
}

// GetApplications will get all applications in the platform.
func (r *Registry) GetApplications(noarg int, reply *[]models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	return db.Find(reply).Error
}

// FindApplication will find product by specified ID
func (r *Registry) FindApplication(id string, reply *models.Application) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Application:%v", id)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		app := cacheValue.(*models.Application)
		setApplication(reply, app)
	} else {
		err = db.First(reply, "record_id = ?", id).Error
		if err != nil {
			return err
		}
		var storage models.Application
		storage = *reply
		cache.Set(cacheKey, &storage)
	}

	return nil
}

//AppAuth will find the application with given args
func (r *Registry) FindApplicationByAppKey(args *rpcs.ArgsAppAuth, reply *models.Application) error {

	db, err := getDB()
	if err != nil {
		return err
	}

	a := &models.Application{}
	a.AppKey = args.AppKey
	a.SecretKey = args.Secretkey

	err = db.Where(a).First(&reply).Error

	if err != nil {
		return err
	}

	return nil
}

// RegisterDevice try to register a device to our platform.
// if the device has already been registered,
// the registration will success return the registered device before.
func (r *Registry) RegisterDevice(args *rpcs.ArgsDeviceRegister, reply *models.Device) error {

	db, err := getDB()
	if err != nil {
		return err
	}

	product := &models.Product{}
	err = r.ValidateProduct(args.ProductKey, product)
	if err != nil {
		return err
	}
	if db.Where(&models.Device{DeviceIdentifier: args.DeviceCode}).First(reply).RecordNotFound() {
		// device is not registered yet.
		reply.RecordId = guid.S()
		reply.ProductID = product.RecordId
		reply.DeviceIdentifier = args.DeviceCode
		reply.DeviceName = product.ProductName // product name as default device name.
		reply.DeviceDescription = product.ProductDescription
		reply.DeviceVersion = args.DeviceVersion
		reply.VendorID = product.VendorID
		reply.ModuleName = args.ModuleName
		err = db.Save(reply).Error
		if err != nil {
			return err
		}
		// generate a random device key with hex encoding.
		reply.DeviceKey, err = r.keygen.GenRandomKey(reply.RecordId)
		if err != nil {
			return err
		}
		// generate a random password with base64 encoding.
		reply.DeviceSecret, err = generator.GenRandomPassword()
		if err != nil {
			return err
		}

		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	} else {

		//delete cache
		cache := getCache()
		cacheKey := fmt.Sprintf("Device:%v", args.DeviceCode)
		if _, ok := cache.Get(cacheKey); ok {
			cache.Delete(cacheKey)
		}

		// device has aleady been saved. just update version info.
		reply.DeviceVersion = args.DeviceVersion
		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	}
	return nil
}

// FindDeviceByIdentifier will find the device by indentifier
func (r *Registry) FindDeviceByIdentifier(identifier string, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	cache := getCache()
	cacheKey := fmt.Sprintf("Device:%v", identifier)
	if cacheValue, ok := cache.Get(cacheKey); ok {
		device := cacheValue.(*models.Device)
		setDevice(reply, device)
	} else {
		err = db.Where(&models.Device{
			DeviceIdentifier: identifier,
		}).First(reply).Error
		if err != nil {
			return err
		}
		cache.Set(cacheKey, reply)
	}

	return nil
}

// FindDeviceByIdentifier2 will find the device by indentifier
func (r *Registry) FindDeviceByIdentifier2(identifier string, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	err = db.Where(&models.Device{
		DeviceIdentifier: identifier,
	}).First(reply).Error
	if err != nil {
		return err
	}

	return nil
}

func (r *Registry) FindDeviceById(args uint64, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	cache := getCache()
	key := fmt.Sprintf("Device:%v", args)
	if v, ok := cache.Get(key); ok {
		device := v.(*models.Device)
		setDevice(reply, device)
	} else {
		d := &models.Device{}
		d.ID = uint(args)
		err = db.Where(d).First(reply).Error
		if err != nil {
			return err
		}
		cache.Set(key, reply)
	}

	return nil
}

// FindDeviceById will find the device with given id
func (r *Registry) FindDeviceByRecordId(args *rpcs.ArgsDeviceAuth, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	cache := getCache()
	key := fmt.Sprintf("Device:%v", args.DeviceID)
	if v, ok := cache.Get(key); ok {
		device := v.(*models.Device)
		setDevice(reply, device)
	} else {
		d := &models.Device{}
		d.RecordId = args.DeviceID
		err = db.Where(d).First(reply).Error
		if err != nil {
			return err
		}
		var storage models.Device
		storage = *reply
		cache.Set(key, &storage)
	}

	return nil
}

// ValidateDevice will validate a device key and return the model if success.
func (r *Registry) ValidateDevice(key string, device *models.Device) error {
	id, err := r.keygen.DecodeIDFromRandomKey(key)
	if err != nil {
		return err
	}
	args := rpcs.ArgsDeviceAuth{
		DeviceID: id,
	}
	err = r.FindDeviceByRecordId(&args, device)
	if err != nil {
		return err
	}

	if device.DeviceKey != key {
		return errors.New("device key not match.")
	}

	return nil
}

// UpdateDeviceInfo will update a device info by identifier
func (r *Registry) UpdateDeviceInfo(args *rpcs.ArgsDeviceUpdate, reply *models.Device) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	err = r.FindDeviceByIdentifier(args.DeviceIdentifier, reply)
	if err != nil {
		return err
	}

	//delete cache
	cache := getCache()
	cacheKey := fmt.Sprintf("Device:%v", args.DeviceIdentifier)
	if _, ok := cache.Get(cacheKey); ok {
		cache.Delete(cacheKey)
	}

	reply.DeviceName = args.DeviceName
	reply.DeviceDescription = args.DeviceDescription

	err = db.Save(reply).Error
	if err != nil {
		return err
	}

	return nil
}

// CreateRule create a new rule with specified parameters.
func (r *Registry) CreateRule(args *models.Rule, reply *rpcs.ReplyEmptyResult) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	return db.Save(args).Error
}

// QueryRules queries rules by trigger and rule type.
func (r *Registry) QueryRules(args *models.Rule, reply *[]models.Rule) error {
	db, err := getDB()
	if err != nil {
		return err
	}

	err = db.Where(args).Find(reply).Error
	if err != nil {
		return err
	}

	return nil
}

// CreateEvent create event
func (r *Registry) CreateEvent(args *models.Event, reply *rpcs.ReplyEmptyResult) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	return db.Save(args).Error
}

// CreateDeviceNetConfig 创建设备配网信息
func (r *Registry) CreateDeviceNetConfig(args *models.DeviceNetConfig, reply *rpcs.ReplyEmptyResult) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	return db.Save(args).Error
}

// CheckDeviceNetConfig 检验设备配网是否成功
func (r *Registry) CheckDeviceNetConfig(args *rpcs.ArgsCheckDeviceNetConfig, reply *rpcs.ReplyCheckDeviceNetConfig) error {
	if db, err := getDB(); err == nil {
		var count int
		err = db.Table("device_net_config").Where("device_identifier = ? and md5 = ?", args.DeviceCode, args.Md5).Count(&count).Error
		if err != nil {
			return err
		}
		if count > 0 {
			reply.Result = true
		} else {
			reply.Result = false
		}
	} else {
		return err
	}
	return nil
}

// UpdateSubDevice  update
func (r *Registry) UpdateSubDevice(args *rpcs.SubDeviceArgs, reply *models.SubDevice) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	if db.Where(&models.SubDevice{SubDeviceID: args.SubDeviceId}).First(reply).RecordNotFound() {
		reply.RecordId = guid.S()
		reply.DeviceCode = args.DeviceCode
		reply.SubDeviceID = args.SubDeviceId
		reply.Status = args.Status

		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	} else {
		// sub_device has already been saved. just update status .
		reply.Status = args.Status
		err = db.Save(reply).Error
		if err != nil {
			return err
		}
	}
	return nil
}

// UpdateDeviceVersion 更新设备版本号
func (r *Registry) UpdateDeviceVersion(args *rpcs.ArgsUpdateDeviceVersion, reply *rpcs.ReplyEmptyResult) error {
	db, err := getDB()
	if err != nil {
		return err
	}
	fReply := &models.Device{}
	err = r.FindDeviceByIdentifier(args.DeviceId, fReply)
	if err != nil {
		return err
	}
	fReply.DeviceVersion = args.Version
	err = db.Save(fReply).Error
	if err != nil {
		return err
	}
	return nil
}