forked from wasabee-project/Wasabee-Server
-
Notifications
You must be signed in to change notification settings - Fork 0
/
model_simple.go
153 lines (133 loc) · 4.32 KB
/
model_simple.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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
package wasabee
import (
"database/sql"
"encoding/hex"
"fmt"
"strings"
"time"
"golang.org/x/crypto/scrypt"
"crypto/sha256"
)
// SimpleDocument specifies the content and metadata of a simple-style (qbin style) draw.
type SimpleDocument struct {
ID string
Content string
Upload time.Time
Expiration time.Time
Views int
}
// Store a document object in the database.
func (document *SimpleDocument) Store() error {
// Generate a name that doesn't exist yet
name, err := GenerateSafeName()
if err != nil {
Log.Errorf("GenerateSafeName: %s", err)
return err
}
document.ID = name
// Round the timestamps on the object. Won't affect the database, but we want consistency.
document.Upload = time.Now().Round(time.Second)
document.Expiration = document.Expiration.Round(time.Second)
// Normalize new lines
document.Content = strings.Trim(strings.Replace(strings.Replace(document.Content, "\r\n", "\n", -1), "\r", "\n", -1), "\n") + "\n"
// Don't accept binary files
if strings.Contains(document.Content, "\x00") {
err = fmt.Errorf("file contains 0x00 bytes")
Log.Debug(err)
return err
}
var expiration interface{}
if (document.Expiration != time.Time{}) {
expiration = document.Expiration.UTC().Format("2006-01-02 15:04:05")
}
// Server-Side Encryption
key, err := scrypt.Key([]byte(document.ID), []byte(document.Upload.UTC().Format("2006-01-02 15:04:05")), 16384, 8, 1, 24)
if err != nil {
Log.Errorf("Invalid script parameters: %s", err)
}
data, err := encrypt([]byte(document.Content), key)
if err != nil {
Log.Errorf("AES error: %s", err)
return err
}
databaseID := sha256.Sum256([]byte(document.ID))
// Write the document to the database
_, err = db.Exec(
"INSERT INTO document (id, content, upload, expiration, views) VALUES (?, ?, ?, ?, 0)",
hex.EncodeToString(databaseID[:]),
string(data),
document.Upload.UTC().Format("2006-01-02 15:04:05"), // don't use NOW() since this is used in the key...
expiration)
if err != nil {
Log.Error(err)
return err
}
return nil
}
// Request a document from the database by its ID.
func Request(id string) (SimpleDocument, error) {
doc := SimpleDocument{ID: id}
var views int
var upload, expiration sql.NullString
databaseID := sha256.Sum256([]byte(id))
err := db.QueryRow("SELECT content, upload, expiration, views FROM document WHERE id = ?", hex.EncodeToString(databaseID[:])).
Scan(&doc.Content, &upload, &expiration, &views)
if err != nil {
if err != sql.ErrNoRows {
Log.Warningf("Error retrieving document: %s", err)
}
return SimpleDocument{}, err
}
_, err = db.Exec("UPDATE document SET views = views + 1 WHERE id = ?", hex.EncodeToString(databaseID[:]))
if err != nil {
Log.Error("unable to update document view count")
}
doc.Views = views
doc.Upload, _ = time.Parse("2006-01-02 15:04:05", upload.String)
key, err := scrypt.Key([]byte(id), []byte(doc.Upload.UTC().Format("2006-01-02 15:04:05")), 16384, 8, 1, 24)
if err != nil {
Log.Errorf("Invalid script parameters: %s", err)
return SimpleDocument{}, err
}
data, err := decrypt([]byte(doc.Content), key)
if err != nil && !(err.Error() == "cipher: message authentication failed" && !strings.Contains(doc.Content, "\000")) {
Log.Errorf("AES error: %s", err)
return SimpleDocument{}, err
} else if err == nil {
doc.Content = string(data)
}
if expiration.Valid {
doc.Expiration, err = time.Parse("2006-01-02 15:04:05", expiration.String)
if doc.Expiration.Before(time.Unix(0, 1)) {
if doc.Views > 0 {
// Volatile document
_, err = db.Exec("DELETE LOW_PRIORITY FROM document WHERE id = ?", hex.EncodeToString(databaseID[:]))
if err != nil {
Log.Errorf("couldn't delete volatile document: %s", err)
}
}
} else {
if err != nil {
return SimpleDocument{}, err
}
if doc.Expiration.Before(time.Now()) {
err = fmt.Errorf("the document has expired")
return SimpleDocument{}, err
}
}
}
return doc, nil
}
func simpleDocClean() {
// do it this way to get RowsAffected
stmt, _ := db.Prepare("DELETE LOW_PRIORITY FROM document WHERE expiration < CURRENT_TIMESTAMP AND expiration > FROM_UNIXTIME(0)")
result, err := stmt.Exec()
if err != nil {
Log.Errorf("couldn't execute cleanup statement: %s", err)
} else {
n, err := result.RowsAffected()
if err == nil && n > 0 {
Log.Debugf("cleaned up %d documents", n)
}
}
}