oc-discovery/models/storage.go
2023-03-08 16:48:36 +01:00

207 lines
5.3 KiB
Go

package models
import (
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"strings"
"github.com/beego/beego/logs"
"github.com/go-redis/redis"
"github.com/tidwall/gjson"
)
type Storage struct {
ZincUrl string
ZincLogin string
ZincPassword string
RedisUrl string
RedisPassword string
}
func LoadPeersJson(filename string) []Peer {
var peers []Peer
content, err := os.ReadFile("./peers.json")
if err != nil {
logs.Error("Error when opening file: ", err)
}
err = json.Unmarshal(content, &peers)
if err != nil {
logs.Error("Error during Unmarshal(): ", err)
}
return peers
}
func (s *Storage) ImportData(peers []Peer) error {
rdb := redis.NewClient(&redis.Options{
Addr: s.RedisUrl,
Password: s.RedisPassword, // no password set
DB: 0, // use default DB
})
var indexedPeers []map[string]interface{}
for _, p := range peers {
// Creating data block for indexing
indexedPeer := make(map[string]interface{})
indexedPeer["_id"] = p.PeerId
indexedPeer["name"] = p.Name
indexedPeer["keywords"] = p.Keywords
indexedPeer["name"] = p.Name
indexedPeer["entityname"] = p.EntityName
indexedPeer["entitytype"] = p.EntityType
indexedPeer["activity"] = p.Activity
indexedPeer["address"] = p.Address
indexedPeer["postcode"] = p.Postcode
indexedPeer["city"] = p.City
indexedPeer["country"] = p.Country
indexedPeer["description"] = p.Description
indexedPeer["apiurl"] = p.ApiUrl
indexedPeer["website"] = p.Website
indexedPeers = append(indexedPeers, indexedPeer)
// Adding peer to Redis (fast retieval and status updates)
jsonp, err := json.Marshal(p)
if err != nil {
return err
}
err = rdb.Set("peer:"+p.PeerId, jsonp, 0).Err()
if err != nil {
return err
}
}
bulk := map[string]interface{}{"index": "peers", "records": indexedPeers}
raw, err := json.Marshal(bulk)
if err != nil {
return err
}
req, err := http.NewRequest("POST", s.ZincUrl+"/api/_bulkv2", strings.NewReader(string(raw)))
if err != nil {
return err
}
req.SetBasicAuth(s.ZincLogin, s.ZincPassword)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
log.Println(resp.StatusCode)
body, err := io.ReadAll(resp.Body)
if err != nil {
return err
}
fmt.Println(string(body))
return nil
}
func (s *Storage) FindPeers(queryString string) ([]Peer, error) {
var peers []Peer
query := `{
"search_type": "match",
"query":
{
"term": "` + queryString + `",
"start_time": "2020-06-02T14:28:31.894Z",
"end_time": "2029-12-02T15:28:31.894Z"
},
"from": 0,
"max_results": 100,
"_source": []
}`
req, err := http.NewRequest("POST", s.ZincUrl+"/api/peers/_search", strings.NewReader(query))
if err != nil {
log.Fatal(err)
}
req.SetBasicAuth(s.ZincLogin, s.ZincPassword)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36")
resp, err := http.DefaultClient.Do(req)
if err != nil {
log.Fatal(err)
}
defer resp.Body.Close()
log.Println(resp.StatusCode)
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
value := gjson.Get(string(body), "hits.hits")
rdb := redis.NewClient(&redis.Options{
Addr: s.RedisUrl,
Password: s.RedisPassword, // no password set
DB: 0, // use default DB
})
for _, v := range value.Array() {
peerBytes, err := rdb.Get("peer:" + v.Get("_id").Str).Bytes()
if err != nil {
logs.Error(err)
} else {
var p Peer
err = json.Unmarshal(peerBytes, &p)
if err != nil {
return nil, err
}
peers = append(peers, p)
}
}
return peers, nil
}
func (s *Storage) GetPeer(uid string) (*Peer, error) {
var peer Peer
rdb := redis.NewClient(&redis.Options{
Addr: s.RedisUrl,
Password: s.RedisPassword, // no password set
DB: 0, // use default DB
})
peerBytes, err := rdb.Get("peer:" + uid).Bytes()
if err != nil {
return nil, err
} else {
err = json.Unmarshal(peerBytes, &peer)
if err != nil {
return nil, err
}
return &peer, nil
}
}
func (s *Storage) DeletePeer(uid string) error {
// Removing from Redis
rdb := redis.NewClient(&redis.Options{
Addr: s.RedisUrl,
Password: s.RedisPassword, // no password set
DB: 0, // use default DB
})
err := rdb.Unlink("peer:" + uid).Err()
if err != nil {
return err
}
// Removing from Index
req, err := http.NewRequest("DELETE", s.ZincUrl+"/api/peers/_doc"+uid, nil)
if err != nil {
return err
}
req.SetBasicAuth(s.ZincLogin, s.ZincPassword)
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
log.Println(resp.StatusCode)
body, err := io.ReadAll(resp.Body)
if err != nil {
return err
}
fmt.Println(string(body))
return nil
}