Skip to content

Commit 376628d

Browse files
kmrozkmroz
authored
Merge pull request #44 from mrozitron/cmd-families
enumerate c2 families
2 parents a87abc6 + 89c42cd commit 376628d

File tree

3 files changed

+235
-22
lines changed

3 files changed

+235
-22
lines changed

‎cmd/get/get.go‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
// Package get implements the get command.
2+
package get
3+
4+
import (
5+
"flag"
6+
"fmt"
7+
"os"
8+
"strings"
9+
"text/tabwriter"
10+
"time"
11+
12+
"github.com/alphasoc/flightsim/wisdom"
13+
)
14+
15+
var Version string
16+
17+
var usage = `usage: flightsim get [flags] element:category
18+
19+
Available elements:
20+
21+
%s
22+
23+
Available categories:
24+
25+
%s
26+
27+
Available flags:
28+
`
29+
30+
// printWelcome prints a basic welcome banner.
31+
func printWelcome() {
32+
fmt.Printf(`
33+
AlphaSOC Network Flight Simulator™ %s (https://github.com/alphasoc/flightsim)
34+
The current time is %s
35+
`, Version, time.Now().Format("02-Jan-06 15:04:05"))
36+
}
37+
38+
// printGoodbye prints a parting message.
39+
func printGoodbye() {
40+
fmt.Printf("\nAll done!\n")
41+
}
42+
43+
// printMsg prints msg, decorated with an info string and the current date/time.
44+
func printMsg(info string, msg string) {
45+
if msg == "" {
46+
return
47+
}
48+
fmt.Printf("%s [%s] %s\n", time.Now().Format("15:04:05"), info, msg)
49+
}
50+
51+
// We only know how to get families for now.
52+
var supportedElementsMap = map[string]bool{
53+
"families": true,
54+
}
55+
56+
// supportedElements returns a slice of fetchable elements.
57+
func supportedElements() []string {
58+
var elements []string
59+
for e := range supportedElementsMap {
60+
elements = append(elements, e)
61+
}
62+
return elements
63+
}
64+
65+
// Taken from open-wisdom/server/entries/entries.go.
66+
// TODO: would be good to fetch these also.
67+
var supportedCategoriesMap = map[string]bool{
68+
"c2": true,
69+
}
70+
71+
// supportedCategories returns a slice of fetchable categories.
72+
func supportedCategories() []string {
73+
var categories []string
74+
for e := range supportedCategoriesMap {
75+
categories = append(categories, e)
76+
}
77+
return categories
78+
}
79+
80+
// computeFormatStr returns a format string to be used with column printing.
81+
func computeFormatStr(cols int) string {
82+
fmtStr := "\n"
83+
for i := 0; i < cols; i++ {
84+
fmtStr += "%v\t"
85+
}
86+
return fmtStr
87+
}
88+
89+
// printCol prints elements in cols columns.
90+
func printCol(elements []string, cols int) {
91+
w := new(tabwriter.Writer)
92+
// Min width, tab width, padding, pad char, flags.
93+
w.Init(os.Stdout, 8, 8, 1, '\t', 0)
94+
// Compute format string.
95+
fmtStr := computeFormatStr(cols)
96+
// Convert elements ([]string) to []interface{}.
97+
elementsToPrint := make([]interface{}, len(elements))
98+
for i, v := range elements {
99+
elementsToPrint[i] = v
100+
}
101+
// Print.
102+
i := 0
103+
lenElementsToPrint := len(elementsToPrint)
104+
for leftToPrint := lenElementsToPrint; leftToPrint > 0; {
105+
// We don't have enough elements left to print to satisfy the format string,
106+
// or the re-slice. Thus, reset *cols and recompute the format string.
107+
if leftToPrint < cols {
108+
cols = leftToPrint
109+
fmtStr = computeFormatStr(cols)
110+
}
111+
fmt.Fprintf(w, fmtStr, elementsToPrint[i:i+cols]...)
112+
// Move by number of columns.
113+
i += cols
114+
leftToPrint = lenElementsToPrint - i
115+
}
116+
// Append a blank line.
117+
fmt.Fprintf(w, "%v", "\n\n")
118+
w.Flush()
119+
}
120+
121+
// RunCmd runs the 'get' command and returns an error.
122+
func RunCmd(args []string) error {
123+
printWelcome()
124+
// Mirrors look of run command.
125+
fmt.Println("")
126+
cmdLine := flag.NewFlagSet("get", flag.ExitOnError)
127+
// TODO: replace cols with -format (issue #45).
128+
// cols := cmdLine.Int("cols", 0, "print elements in number of columns")
129+
cmdLine.Usage = func() {
130+
fmt.Fprintf(cmdLine.Output(), usage, strings.Join(supportedElements(), ", "), strings.Join(supportedCategories(), ", "))
131+
cmdLine.PrintDefaults()
132+
}
133+
cmdLine.Parse(args)
134+
// Next arg should be element:category (ie. families:c2)
135+
toGet := cmdLine.Arg(0)
136+
if len(toGet) == 0 {
137+
return fmt.Errorf("nothing to get")
138+
}
139+
toGetArr := strings.Split(cmdLine.Arg(0), ":")
140+
if len(toGetArr) != 2 {
141+
return fmt.Errorf("unable to get '%v': invalid format", toGet)
142+
}
143+
elem := toGetArr[0]
144+
cat := toGetArr[1]
145+
// infoTag == "element:category" (ie. families:c2). Mirrors the run command.
146+
infoTag := cmdLine.Arg(0)
147+
var elements []string
148+
var err error
149+
switch elem {
150+
case "families":
151+
printMsg(infoTag, fmt.Sprintf("Fetching %v %v", cat, elem))
152+
elements, err = wisdom.Families(cat)
153+
if err != nil {
154+
return err
155+
}
156+
default:
157+
return fmt.Errorf("unable to get '%v': unsupported element '%v'", toGet, elem)
158+
}
159+
// TODO: Leave quotes in for later cut/paste?
160+
for i := 0; i < len(elements); i++ {
161+
elements[i] = strings.Trim(elements[i], "\"")
162+
}
163+
// Default is to print in a single line. Otherwise, column print.
164+
// if *cols <= 0 {
165+
// printMsg(infoTag, strings.Join(elements, ", "))
166+
// } else {
167+
// printCol(elements, *cols)
168+
// }
169+
printMsg(infoTag, strings.Join(elements, ", "))
170+
171+
printMsg(infoTag, fmt.Sprintf("Fetched %v %v %v", len(elements), cat, elem))
172+
printGoodbye()
173+
return nil
174+
}

‎main.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"os"
77

8+
"github.com/alphasoc/flightsim/cmd/get"
89
"github.com/alphasoc/flightsim/cmd/run"
910
)
1011

@@ -20,10 +21,12 @@ Usage:
2021
flightsim <command> [arguments]
2122
2223
Available commands:
24+
get Get a list of elements (ie. families) of a certain category (ie. c2)
2325
run Run all modules, or a particular module
2426
version Prints the version number
2527
2628
Cheatsheet:
29+
flightsim get families:c2 Get a list of all c2 families
2730
flightsim run Run all the modules
2831
flightsim run c2 Simulate C2 traffic
2932
flightsim run c2:trickbot Simulate C2 traffic for the TrickBot family
@@ -48,6 +51,9 @@ func main() {
4851
var err error
4952

5053
switch cmd {
54+
case "get":
55+
get.Version = Version
56+
err = get.RunCmd(args)
5157
case "run":
5258
run.Version = Version
5359
err = run.RunCmd(args)

‎wisdom/wisdom.go‎

Lines changed: 55 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -56,31 +56,10 @@ func (h *WisdomHosts) Hosts(scope string, size int) ([]string, error) {
5656
reqURL.Query().Set("family", h.Family)
5757
}
5858

59-
client := &http.Client{
60-
Timeout: 5 * time.Second,
61-
}
62-
63-
var resp *http.Response
64-
for n := 0; n < 3; n++ {
65-
resp, err = client.Get(reqURL.String())
66-
if err == nil {
67-
break
68-
}
69-
}
59+
b, err := query(reqURL)
7060
if err != nil {
7161
return nil, err
7262
}
73-
defer resp.Body.Close()
74-
75-
if c := resp.StatusCode; c != http.StatusOK {
76-
b, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 2048))
77-
return nil, fmt.Errorf("api.open.wisdom.alphasoc.net said: %d: %s", c, b)
78-
}
79-
80-
b, err := ioutil.ReadAll(resp.Body)
81-
if err != nil {
82-
return nil, errors.Wrapf(err, "api.open.wisdom.alphasoc.net read body error")
83-
}
8463

8564
var parsed struct {
8665
Items []struct {
@@ -122,3 +101,57 @@ func (h *WisdomHosts) Hosts(scope string, size int) ([]string, error) {
122101

123102
return hosts, nil
124103
}
104+
105+
// Families queries the wisdom families API, returning families for a given category
106+
// as a slice of strings, along with an error.
107+
func Families(category string) ([]string, error) {
108+
reqURL, err := url.Parse("https://api.open.wisdom.alphasoc.net/v2/families")
109+
if err != nil {
110+
return nil, err
111+
}
112+
q := reqURL.Query()
113+
q.Set("category", category)
114+
reqURL.RawQuery = q.Encode()
115+
b, err := query(reqURL)
116+
if err != nil {
117+
return nil, err
118+
}
119+
var parsed struct {
120+
Families []string
121+
}
122+
if err := json.Unmarshal(b, &parsed); err != nil {
123+
return nil, errors.Wrapf(err, "api.open.wisdom.alphasoc.net parse body error")
124+
}
125+
return parsed.Families, nil
126+
}
127+
128+
// query carries out a wisdom query, returning a byte slice and an error.
129+
func query(reqURL *url.URL) ([]byte, error) {
130+
client := &http.Client{
131+
Timeout: 5 * time.Second,
132+
}
133+
134+
var resp *http.Response
135+
var err error
136+
for n := 0; n < 3; n++ {
137+
resp, err = client.Get(reqURL.String())
138+
if err == nil {
139+
break
140+
}
141+
}
142+
if err != nil {
143+
return nil, err
144+
}
145+
defer resp.Body.Close()
146+
147+
if c := resp.StatusCode; c != http.StatusOK {
148+
b, _ := ioutil.ReadAll(io.LimitReader(resp.Body, 2048))
149+
return nil, fmt.Errorf("api.open.wisdom.alphasoc.net said: %d: %s", c, b)
150+
}
151+
152+
b, err := ioutil.ReadAll(resp.Body)
153+
if err != nil {
154+
return nil, errors.Wrapf(err, "api.open.wisdom.alphasoc.net read body error")
155+
}
156+
return b, nil
157+
}

0 commit comments

Comments
 (0)