set_tenant_restricted.md

package main

import (
	"fmt"
	"os"

	"github.com/hyperhq/hyper-api-router/pkg/auth"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"k8s.io/kubernetes/pkg/util"
)

type TenantRestricted struct {
	DB string

	Fire bool
}

func (s *TenantRestricted) AddFlags(fs *pflag.FlagSet) {
	fs.BoolVar(&s.Fire, "fire", s.Fire, "really do the operation")
	fs.StringVar(&s.DB, "db", s.DB, "url of the source db, default 127.0.0.1")
}

func (s *TenantRestricted) Run(_ []string) error {
	var (
		mongo *mgo.Session
		err   error
	)

	mongo, err = mgo.Dial(s.DB)
	if err != nil {
		glog.Errorf(err.Error())
		return err
	}
	defer mongo.Close()

	miter := mongo.Copy()
	defer miter.Close()

	var tenant auth.TenantItem
	iter := miter.DB("hypernetes").C("tenant").Find(bson.M{}).Iter()
	for iter.Next(&tenant) {
		glog.Infof("Operate tenant %s", tenant.TenantID)
		if tenant.Restricted {
			continue
		}

		count, err := mongo.DB("hypernetes").C("credentials").Find(bson.M{
			"tenantId": tenant.TenantID,
			"system":   false,
		}).Count()
		if err != nil {
			glog.Errorf(err.Error())
			return err
		}
		if count > 0 {
			glog.Infof("%s has %d credentials, pass", tenant.TenantID, count)
			continue
		}
		glog.Infof("%s should be set", tenant.TenantID)

		if !s.Fire {
			continue
		}

		err = mongo.DB("hypernetes").C("tenant").Update(bson.M{
			"tenantid": tenant.TenantID,
		}, bson.M{
			"$set": bson.M{
				"restricted": true,
			},
		})
		if err != nil {
			glog.Errorf(err.Error())
			return err
		}

		glog.Infof("successfully updated tenant %s to false", tenant.TenantID)
	}

	if err = iter.Err(); err != nil {
		glog.Errorf(err.Error())
		return err
	}
	iter.Close()

	return nil
}

func NewTenantRestricted() *TenantRestricted {
	return &TenantRestricted{
		DB: "127.0.0.1",
	}
}

func main() {
	t := NewTenantRestricted()
	t.AddFlags(pflag.CommandLine)

	util.InitFlags()
	util.InitLogs()

	if err := t.Run(pflag.CommandLine.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

}