forked from engineering/cloudservices
109 lines
3.0 KiB
Go
109 lines
3.0 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/go-redis/redis/v7"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
// RedisClient the Redis client
|
|
var RedisClient *redis.Client
|
|
|
|
// ConfigStruct the configuration struct
|
|
type ConfigStruct struct {
|
|
MongoDB string `json:"mongoURL"`
|
|
}
|
|
|
|
// Account represents an user's account.
|
|
type Account struct {
|
|
Username string `json:"username"`
|
|
UserID string `json:"userId"`
|
|
EmailAddress string `json:"emailAddress"`
|
|
CreatedBy string `json:"createdBy"`
|
|
CreatedAt time.Time `json:"createdAt"`
|
|
Locked bool `json:"locked"`
|
|
}
|
|
|
|
// HandleError This function panics if an error exists.
|
|
func HandleError(e error, serv int) {
|
|
if e != nil && serv == 1 {
|
|
panic(e)
|
|
} else if e != nil && serv == 0 {
|
|
fmt.Println(e)
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
fmt.Println("Starting CSD-GO-STR Daemon [GO]")
|
|
config := &ConfigStruct{}
|
|
file, err := ioutil.ReadFile("../config.json")
|
|
HandleError(err, 1)
|
|
err = json.Unmarshal(file, &config)
|
|
|
|
RedisClient = redis.NewClient(&redis.Options{
|
|
Addr: "localhost:6379",
|
|
Password: "",
|
|
DB: 0,
|
|
})
|
|
|
|
_, err = RedisClient.Ping().Result()
|
|
fmt.Printf("Connected to Redis [GO]\n")
|
|
HandleError(err, 1)
|
|
|
|
/*for {
|
|
fmt.Printf("Calling handler func [GO]\n")
|
|
if status == false {
|
|
handler(*config)
|
|
time.Sleep(1000000 * time.Millisecond)
|
|
}
|
|
}*/
|
|
handler(*config)
|
|
}
|
|
|
|
func handler(config ConfigStruct) {
|
|
mongoClient, err := mongo.NewClient(options.Client().ApplyURI(config.MongoDB))
|
|
HandleError(err, 1)
|
|
ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
|
|
err = mongoClient.Connect(ctx)
|
|
HandleError(err, 1)
|
|
err = mongoClient.Ping(context.TODO(), nil)
|
|
fmt.Printf("Connected to MongoDB [GO]\n")
|
|
HandleError(err, 1)
|
|
collection := mongoClient.Database("cloudservices").Collection("accounts")
|
|
cur, err := collection.Find(context.TODO(), bson.D{})
|
|
HandleError(err, 0)
|
|
|
|
for cur.Next(context.TODO()) {
|
|
checkAccountSizeAndUpdate(cur.Current.Lookup("username").String(), cur.Current.Lookup("id").String())
|
|
fmt.Printf("Checking account information for %s\n", cur.Current.Lookup("username").String())
|
|
time.Sleep(10 * time.Second)
|
|
}
|
|
err = mongoClient.Disconnect(ctx)
|
|
HandleError(err, 1)
|
|
fmt.Println("CSD-GO-STR finished, exiting with code 0. [GO]")
|
|
os.Exit(0)
|
|
}
|
|
|
|
func checkAccountSizeAndUpdate(username string, id string) {
|
|
var size float64 = 0
|
|
var userHomeDirectory string = strings.Replace(strings.Join([]string{"/home/", string(username)}, ""), "\"", "", -1)
|
|
usernameFormat := strings.Replace(username, "\"", "", -1)
|
|
sizeHome := DirSize(&userHomeDirectory)
|
|
size += sizeHome
|
|
status := RedisClient.Set("storage"+"-"+usernameFormat, size, 0)
|
|
fmt.Printf("Redis Status: %s [GO]\n", status.Name())
|
|
if status.Err() != nil {
|
|
fmt.Println(status.Err())
|
|
}
|
|
fmt.Printf("Set Call | Username: %v, ID: %v | Bytes: %f [GO]\n", string(username), string(id), size)
|
|
}
|