/
mythic_classes.go
122 lines (106 loc) · 2.92 KB
/
mythic_classes.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
package gomythic
import (
"encoding/json"
"github.com/gorilla/websocket"
"time"
)
type Mythic struct {
Username string `json:"username"`
Password string `json:"password"`
APIToken string `json:"apitoken"`
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ServerIP string `json:"server_ip"`
ServerPort int `json:"server_port"`
SSL bool `json:"ssl"`
HTTP string `json:"-"`
WS string `json:"-"`
GlobalTimeout int `json:"global_timeout"`
ScriptingVersion string `json:"scripting_version"`
CurrentOperationID int `json:"current_operation_id"`
Schema string `json:"schema"`
}
type MythicWebSocketHandler struct {
Conn *websocket.Conn
timeout time.Duration
}
func (m *Mythic) String() string {
data, _ := json.MarshalIndent(m, "", " ")
return string(data)
}
type MythicStatus string
const (
Error MythicStatus = "error"
Completed MythicStatus = "completed"
Processed MythicStatus = "processed"
Processing MythicStatus = "processing"
Preprocessing MythicStatus = "preprocessing"
Delegating MythicStatus = "delegating"
Submitted MythicStatus = "submitted"
)
func (s MythicStatus) String() string {
return string(s)
}
func (s MythicStatus) Equals(obj interface{}) bool {
if status, ok := obj.(string); ok {
return s.String() == status
} else if status, ok := obj.(MythicStatus); ok {
return s.String() == status.String()
}
return false
}
func (s MythicStatus) GreaterThanOrEqual(obj interface{}) bool {
targetObj := ""
if status, ok := obj.(string); ok {
targetObj = status
} else if status, ok := obj.(MythicStatus); ok {
targetObj = status.String()
}
if targetObj == "" {
return false
}
if targetObj == "delegating" {
targetObj = "delegating"
} else if targetObj == "error" {
targetObj = "error"
}
selfObj := s.String()
if selfObj == "delegating" {
selfObj = "delegating"
}
if selfObj == "error" {
return true
} else if selfObj == "completed" {
return true
}
enumMapping := map[string]int{
"preprocessing": 0,
"submitted": 1,
"delegating": 2,
"processing": 3,
"processed": 4,
"completed": 5,
"error": 6,
}
if _, ok := enumMapping[targetObj]; !ok {
panic("Can't compare status of type: " + targetObj)
} else if _, ok := enumMapping[selfObj]; !ok {
panic("Can't compare status of type: " + selfObj)
}
return enumMapping[selfObj] >= enumMapping[targetObj]
}
type MythicManager struct {
instance *Mythic
}
func NewMythicManager() *MythicManager {
return &MythicManager{}
}
func (mm *MythicManager) GetMythicInstance() *Mythic {
if mm.instance == nil {
mm.instance = &Mythic{} // create a new instance here
}
return mm.instance
}
func (mm *MythicManager) InvalidateMythicInstance() {
mm.instance = nil // set the instance to nil
}