/
potfile.go
104 lines (84 loc) · 2.12 KB
/
potfile.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
package db
import (
"errors"
"github.com/lachlan2k/phatcrack/common/pkg/apitypes"
"gorm.io/gorm"
)
type PotfileEntry struct {
UUIDBaseModel
Hash string `gorm:"index:idx_potfile_hash,type:hash"`
PlaintextHex string
HashType uint
}
type PotfileSearchResult struct {
Entry *PotfileEntry
Hash string
Found bool
}
func (r PotfileSearchResult) ToDTO() apitypes.PotfileSearchResultDTO {
dto := apitypes.PotfileSearchResultDTO{
Hash: r.Hash,
Found: r.Found,
}
if r.Entry != nil {
dto.PlaintextHex = r.Entry.PlaintextHex
dto.HashType = r.Entry.HashType
}
return dto
}
func AddPotfileEntry(newEntry *PotfileEntry) (*PotfileEntry, error) {
err := GetInstance().Transaction(func(tx *gorm.DB) error {
var foundEntry PotfileEntry
err := tx.First(&foundEntry, "hash = ? and hash_type = ?", newEntry.Hash, newEntry.HashType).Error
// If it wasn't found, create it
if errors.Is(err, gorm.ErrRecordNotFound) {
return tx.Create(newEntry).Error
}
if err != nil {
return err
}
// or, if it was found, but is a colission, create it
if foundEntry.PlaintextHex != newEntry.PlaintextHex {
return tx.Create(newEntry).Error
}
return nil
})
if err != nil {
return nil, err
}
return newEntry, nil
}
func SearchPotfile(hashes []string) ([]PotfileSearchResult, error) {
results := make([]PotfileSearchResult, 0)
err := GetInstance().Transaction(func(tx *gorm.DB) error {
for _, loopHashToSearch := range hashes {
hashToSearch := loopHashToSearch
foundEntries := []PotfileEntry{}
err := tx.Where("hash = ?", hashToSearch).Find(&foundEntries).Error
if errors.Is(err, gorm.ErrRecordNotFound) || len(foundEntries) == 0 {
results = append(results, PotfileSearchResult{
Entry: nil,
Hash: hashToSearch,
Found: false,
})
continue
}
if err != nil {
return err
}
for _, loopFoundEntry := range foundEntries {
foundEntry := loopFoundEntry
results = append(results, PotfileSearchResult{
Entry: &foundEntry,
Hash: hashToSearch,
Found: true,
})
}
}
return nil
})
if err != nil {
return nil, err
}
return results, nil
}