Skip to content

Commit 7fcf3fc

Browse files
heiytorheiytor
committed
wip
1 parent 7b0b4a4 commit 7fcf3fc

31 files changed

+1585
-5
lines changed

‎api/store/pg/device.go‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
package pg
2+
3+
import (
4+
"context" //nolint:gosec
5+
"time"
6+
7+
"github.com/shellhub-io/shellhub/api/store"
8+
"github.com/shellhub-io/shellhub/api/store/pg/entity"
9+
"github.com/shellhub-io/shellhub/pkg/clock"
10+
"github.com/shellhub-io/shellhub/pkg/models"
11+
"github.com/uptrace/bun"
12+
"github.com/uptrace/bun/dialect/pgdialect"
13+
)
14+
15+
func (pg *Pg) DeviceCreate(ctx context.Context, device *models.Device) (string, error) {
16+
device.CreatedAt = clock.Now()
17+
18+
e := entity.DeviceFromModel(device)
19+
if _, err := pg.driver.NewInsert().Model(e).Exec(ctx); err != nil {
20+
return "", fromSqlError(err)
21+
}
22+
23+
return e.ID, nil
24+
}
25+
26+
func (pg *Pg) DeviceConflicts(ctx context.Context, target *models.DeviceConflicts) ([]string, bool, error) {
27+
devices := make([]map[string]any, 0)
28+
if err := pg.driver.NewSelect().Model((*entity.Device)(nil)).Column("name").Where("name = ?", target.Name).Scan(ctx, &devices); err != nil {
29+
return nil, false, fromSqlError(err)
30+
}
31+
32+
conflicts := make([]string, 0)
33+
for _, device := range devices {
34+
if device["name"] == target.Name {
35+
conflicts = append(conflicts, "name")
36+
}
37+
}
38+
39+
return conflicts, len(conflicts) > 0, nil
40+
}
41+
42+
func (pg *Pg) DeviceList(ctx context.Context, opts ...store.QueryOption) ([]models.Device, int, error) {
43+
entities := make([]entity.Device, 0)
44+
45+
query := pg.driver.
46+
NewSelect().
47+
Model(&entities).
48+
Column("device.*").
49+
Relation("Namespace").
50+
ColumnExpr(`
51+
CASE
52+
WHEN "device"."disconnected_at" IS NULL AND "device"."seen_at" > ?
53+
THEN true
54+
ELSE false
55+
END AS "online"`,
56+
time.Now().Add(-2*time.Minute),
57+
).
58+
ColumnExpr(`
59+
CASE
60+
WHEN "device"."status" <> 'accepted'
61+
THEN true
62+
ELSE false
63+
END AS "acceptable"`,
64+
)
65+
66+
if err := applyOptions(ctx, query, opts...); err != nil {
67+
return nil, 0, fromSqlError(err)
68+
}
69+
70+
count, err := query.ScanAndCount(ctx)
71+
if err != nil {
72+
return nil, 0, fromSqlError(err)
73+
}
74+
75+
devices := make([]models.Device, len(entities))
76+
for i, e := range entities {
77+
devices[i] = *entity.DeviceToModel(&e)
78+
}
79+
80+
return devices, count, nil
81+
}
82+
83+
func (pg *Pg) DeviceResolve(ctx context.Context, resolver store.DeviceResolver, val string, opts ...store.QueryOption) (*models.Device, error) {
84+
d := new(entity.Device)
85+
86+
query := pg.driver.
87+
NewSelect().
88+
Model(d).
89+
Where("? = ?", bun.Ident("device."+string(resolver)), val).
90+
Column("device.*").
91+
Relation("Namespace").
92+
ColumnExpr(`
93+
CASE
94+
WHEN "device"."disconnected_at" IS NULL AND "device"."seen_at" > ?
95+
THEN true
96+
ELSE false
97+
END AS "online"`,
98+
time.Now().Add(-2*time.Minute),
99+
)
100+
101+
if err := query.Scan(ctx); err != nil {
102+
return nil, fromSqlError(err)
103+
}
104+
105+
return entity.DeviceToModel(d), nil
106+
}
107+
108+
func (pg *Pg) DeviceSave(ctx context.Context, device *models.Device) error {
109+
d := entity.DeviceFromModel(device)
110+
d.UpdatedAt = clock.Now()
111+
_, err := pg.driver.NewUpdate().Model(d).WherePK().Exec(ctx)
112+
113+
return fromSqlError(err)
114+
}
115+
116+
func (pg *Pg) DeviceUpdateSeenAt(ctx context.Context, ids []string, to time.Time) (int64, error) {
117+
r, err := pg.driver.NewUpdate().
118+
Model((*entity.Device)(nil)).
119+
Set("seen_at = ?", to).
120+
Set("disconnected_at = NULL").
121+
TableExpr("(SELECT unnest(?::varchar[]) as id) as _data", pgdialect.Array(ids)).
122+
Where("device.id = _data.id").
123+
Exec(ctx)
124+
if err != nil {
125+
return 0, fromSqlError(err)
126+
}
127+
128+
return r.RowsAffected()
129+
}
130+
131+
func (pg *Pg) DeviceDelete(ctx context.Context, device *models.Device) error {
132+
d := entity.DeviceFromModel(device)
133+
_, err := pg.driver.NewDelete().Model(d).WherePK().Exec(ctx)
134+
135+
return fromSqlError(err)
136+
}

‎api/store/pg/device_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
package pg_test

‎api/store/pg/entity/device.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
package entity
2+
3+
import (
4+
"time"
5+
6+
"github.com/shellhub-io/shellhub/pkg/models"
7+
"github.com/uptrace/bun"
8+
)
9+
10+
type Device struct {
11+
bun.BaseModel `bun:"table:devices"`
12+
13+
ID string `bun:"id,pk"`
14+
NamespaceID string `bun:"namespace_id,pk,type:uuid"`
15+
CreatedAt time.Time `bun:"created_at"`
16+
UpdatedAt time.Time `bun:"updated_at"`
17+
SeenAt time.Time `bun:"seen_at"`
18+
DisconnectedAt time.Time `bun:"disconnected_at,nullzero"`
19+
Online bool `bun:",scanonly"`
20+
Acceptable bool `bun:",scanonly"`
21+
Status string `bun:"status"`
22+
Name string `bun:"name"`
23+
MAC string `bun:"mac"`
24+
PublicKey string `bun:"public_key"`
25+
Identifier string `bun:"identifier"`
26+
PrettyName string `bun:"pretty_name"`
27+
Version string `bun:"version"`
28+
Arch string `bun:"arch"`
29+
Platform string `bun:"platform"`
30+
Longitude float64 `bun:"longitude,type:numeric"`
31+
Latitude float64 `bun:"latitude,type:numeric"`
32+
33+
Namespace *Namespace `bun:"rel:belongs-to,join:namespace_id=id"`
34+
}
35+
36+
func DeviceFromModel(model *models.Device) *Device {
37+
device := &Device{
38+
ID: model.UID,
39+
NamespaceID: model.TenantID,
40+
CreatedAt: model.CreatedAt,
41+
UpdatedAt: time.Time{},
42+
SeenAt: model.LastSeen,
43+
Status: string(model.Status),
44+
Name: model.Name,
45+
PublicKey: model.PublicKey,
46+
}
47+
48+
if model.DisconnectedAt != nil {
49+
device.DisconnectedAt = *model.DisconnectedAt
50+
}
51+
52+
if model.Identity != nil {
53+
device.MAC = model.Identity.MAC
54+
}
55+
56+
if model.Position != nil {
57+
device.Longitude = model.Position.Longitude
58+
device.Latitude = model.Position.Latitude
59+
}
60+
61+
if model.Info != nil {
62+
device.Identifier = model.Info.ID
63+
device.PrettyName = model.Info.PrettyName
64+
device.Version = model.Info.Version
65+
device.Arch = model.Info.Arch
66+
device.Platform = model.Info.Platform
67+
}
68+
69+
return device
70+
}
71+
72+
func DeviceToModel(entity *Device) *models.Device {
73+
device := &models.Device{
74+
UID: entity.ID,
75+
TenantID: entity.NamespaceID,
76+
CreatedAt: entity.CreatedAt,
77+
LastSeen: entity.SeenAt,
78+
Status: models.DeviceStatus(entity.Status),
79+
Name: entity.Name,
80+
PublicKey: entity.PublicKey,
81+
Online: entity.Online,
82+
Acceptable: entity.Acceptable,
83+
Namespace: entity.Namespace.Name,
84+
DisconnectedAt: nil,
85+
RemoteAddr: "",
86+
Tags: []string{},
87+
Position: &models.DevicePosition{
88+
Longitude: entity.Longitude,
89+
Latitude: entity.Latitude,
90+
},
91+
Info: &models.DeviceInfo{
92+
ID: entity.Identifier,
93+
PrettyName: entity.PrettyName,
94+
Version: entity.Version,
95+
Arch: entity.Arch,
96+
Platform: entity.Platform,
97+
},
98+
Identity: &models.DeviceIdentity{
99+
MAC: entity.MAC,
100+
},
101+
}
102+
103+
if !entity.DisconnectedAt.IsZero() {
104+
disconnectedAt := entity.DisconnectedAt
105+
device.DisconnectedAt = &disconnectedAt
106+
}
107+
108+
return device
109+
}

‎api/store/pg/entity/private-key.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package entity
2+
3+
import (
4+
"time"
5+
6+
"github.com/shellhub-io/shellhub/pkg/models"
7+
"github.com/uptrace/bun"
8+
)
9+
10+
type PrivateKey struct {
11+
bun.BaseModel `bun:"table:private_keys"`
12+
13+
Fingerprint string `bun:"fingerprint,pk"`
14+
CreatedAt time.Time `bun:"created_at"`
15+
UpdatedAt time.Time `bun:"updated_at"`
16+
Data []byte `bun:"data,type:bytea"`
17+
}
18+
19+
func PrivateKeyFromModel(model *models.PrivateKey) *PrivateKey {
20+
return &PrivateKey{
21+
Fingerprint: model.Fingerprint,
22+
Data: model.Data,
23+
CreatedAt: model.CreatedAt,
24+
UpdatedAt: time.Time{},
25+
}
26+
}
27+
28+
func PrivateKeyToModel(entity *PrivateKey) *models.PrivateKey {
29+
return &models.PrivateKey{
30+
Fingerprint: entity.Fingerprint,
31+
Data: entity.Data,
32+
CreatedAt: entity.CreatedAt,
33+
}
34+
}

‎api/store/pg/entity/public-key.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
package entity
2+
3+
import (
4+
"time"
5+
6+
"github.com/shellhub-io/shellhub/pkg/models"
7+
"github.com/uptrace/bun"
8+
)
9+
10+
type PublicKey struct {
11+
bun.BaseModel `bun:"table:public_keys"`
12+
13+
ID string `bun:"id,pk"`
14+
Fingerprint string `bun:"fingerprint"`
15+
NamespaceID string `bun:"namespace_id"`
16+
CreatedAt time.Time `bun:"created_at"`
17+
UpdatedAt time.Time `bun:"updated_at"`
18+
Name string `bun:"name"`
19+
Data []byte `bun:"data,type:bytea"`
20+
}
21+
22+
func PublicKeyFromModel(model *models.PublicKey) *PublicKey {
23+
return &PublicKey{
24+
NamespaceID: model.TenantID,
25+
Fingerprint: model.Fingerprint,
26+
CreatedAt: model.CreatedAt,
27+
UpdatedAt: time.Time{},
28+
Name: model.PublicKeyFields.Name,
29+
Data: model.Data,
30+
}
31+
}
32+
33+
func PublicKeyToModel(entity *PublicKey) *models.PublicKey {
34+
return &models.PublicKey{
35+
TenantID: entity.NamespaceID,
36+
Fingerprint: entity.Fingerprint,
37+
Data: entity.Data,
38+
CreatedAt: entity.CreatedAt,
39+
PublicKeyFields: models.PublicKeyFields{
40+
Name: entity.Name,
41+
Username: "",
42+
Filter: models.PublicKeyFilter{
43+
Hostname: "",
44+
Tags: []string{},
45+
},
46+
},
47+
}
48+
}

0 commit comments

Comments
 (0)