Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
Path: blob/dev/pkg/testutils/fuzzplayground/db.go
2070 views
1
package fuzzplayground
2
3
import (
4
"database/sql"
5
"encoding/xml"
6
"fmt"
7
"os"
8
"strconv"
9
"strings"
10
11
_ "github.com/mattn/go-sqlite3"
12
)
13
14
var (
15
db *sql.DB
16
tempDBDir string
17
)
18
19
func init() {
20
dir, err := os.MkdirTemp("", "fuzzplayground-*")
21
if err != nil {
22
panic(err)
23
}
24
tempDBDir = dir
25
26
db, err = sql.Open("sqlite3", fmt.Sprintf("file:%v/test.db?cache=shared&mode=memory", tempDBDir))
27
if err != nil {
28
panic(err)
29
}
30
addDummyUsers(db)
31
addDummyPosts(db)
32
}
33
34
// Cleanup cleans up the temporary database directory
35
func Cleanup() {
36
if db != nil {
37
_ = db.Close()
38
}
39
if tempDBDir != "" {
40
_ = os.RemoveAll(tempDBDir)
41
}
42
}
43
44
type User struct {
45
XMLName xml.Name `xml:"user"`
46
ID int `xml:"id"`
47
Name string `xml:"name"`
48
Age int `xml:"age"`
49
Role string `xml:"role"`
50
}
51
52
func addDummyUsers(db *sql.DB) {
53
_, err := db.Exec("CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER, role TEXT)")
54
if err != nil {
55
panic(err)
56
}
57
_, err = db.Exec("INSERT INTO users (id , name, age, role) VALUES (1,'admin', 30, 'admin')")
58
if err != nil {
59
panic(err)
60
}
61
_, err = db.Exec("INSERT INTO users (id , name, age, role) VALUES (75,'user', 30, 'user')")
62
if err != nil {
63
panic(err)
64
}
65
}
66
67
func patchUnsanitizedUser(db *sql.DB, user User) error {
68
setClauses := ""
69
70
if user.Name != "" {
71
setClauses += "name = '" + user.Name + "', "
72
}
73
if user.Age > 0 {
74
setClauses += "age = " + strconv.Itoa(user.Age) + ", "
75
}
76
if user.Role != "" {
77
setClauses += "role = '" + user.Role + "', "
78
}
79
if setClauses == "" {
80
// No fields to update
81
return nil
82
}
83
setClauses = strings.TrimSuffix(setClauses, ", ")
84
85
query := "UPDATE users SET " + setClauses + " WHERE id = ?"
86
_, err := db.Exec(query, user.ID)
87
if err != nil {
88
return err
89
}
90
return nil
91
}
92
93
func getUnsanitizedUser(db *sql.DB, id string) (User, error) {
94
var user User
95
err := db.QueryRow("SELECT id, name, age, role FROM users WHERE id = "+id).Scan(&user.ID, &user.Name, &user.Age, &user.Role)
96
if err != nil {
97
return user, err
98
}
99
return user, nil
100
}
101
102
type Posts struct {
103
ID int
104
Title string
105
Content string
106
Lang string
107
}
108
109
func addDummyPosts(db *sql.DB) {
110
_, err := db.Exec("CREATE TABLE IF NOT EXISTS posts (id INTEGER PRIMARY KEY, title TEXT, content TEXT, lang TEXT)")
111
if err != nil {
112
panic(err)
113
}
114
// Inserting English dummy posts
115
_, err = db.Exec("INSERT INTO posts (id, title, content, lang) VALUES (1, 'The Joy of Programming', 'Programming is like painting a canvas with logic.', 'en')")
116
if err != nil {
117
panic(err)
118
}
119
_, err = db.Exec("INSERT INTO posts (id, title, content, lang) VALUES (2, 'A Journey Through Code', 'Every line of code tells a story.', 'en')")
120
if err != nil {
121
panic(err)
122
}
123
// Inserting a Spanish dummy post
124
_, err = db.Exec("INSERT INTO posts (id, title, content, lang) VALUES (3, 'La belleza del código', 'Cada función es un poema en un mar de algoritmos.', 'es')")
125
if err != nil {
126
panic(err)
127
}
128
}
129
130
func getUnsanitizedPostsByLang(db *sql.DB, lang string) ([]Posts, error) {
131
var posts []Posts
132
query := "SELECT id, title, content, lang FROM posts WHERE lang = '" + lang + "'"
133
rows, err := db.Query(query)
134
if err != nil {
135
return nil, err
136
}
137
defer func() {
138
_ = rows.Close()
139
}()
140
141
for rows.Next() {
142
var post Posts
143
if err := rows.Scan(&post.ID, &post.Title, &post.Content, &post.Lang); err != nil {
144
return nil, err
145
}
146
posts = append(posts, post)
147
}
148
if err = rows.Err(); err != nil {
149
return nil, err
150
}
151
return posts, nil
152
}
153
154