Add -n option to run command

alphasoc · Apr 30, 2018 · 81ba56d · 81ba56d
1 parent 37d1eae
commit 81ba56d
Show file tree

Hide file tree

Showing 12 changed files with 26 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -49,6 +49,7 @@ Usage:
   flightsim run [c2-dns|c2-ip|dga|hijack|scan|sink|spambot|tor|tunnel] [flags]
 
 Flags:
+  -n,                      number of hosts generated for each simulator (default 10)
       --fast               run simulator fast without sleep intervals
   -h, --help               help for run
   -i, --interface string   network interface to use

diff --git a/cmd/run.go b/cmd/run.go
@@ -16,6 +16,7 @@ import (
 
 var (
 	fast           bool
+	size           int
 	ifaceName      string
 	simulatorNames = []string{"c2-dns", "c2-ip", "dga", "hijack", "scan", "sink", "spambot", "tunnel"}
 )
@@ -35,6 +36,10 @@ func newRunCommand() *cobra.Command {
 				simulatorNames = args
 			}
 
+			if size <= 0 {
+				return fmt.Errorf("n must be positive")
+			}
+
 			extIP, err := utils.ExternalIP(ifaceName)
 			if err != nil {
 				return err
@@ -54,6 +59,7 @@ func newRunCommand() *cobra.Command {
 	}
 
 	cmd.Flags().BoolVar(&fast, "fast", false, "run simulator fast without sleep intervals")
+	cmd.Flags().IntVarP(&size, "", "n", 10, "number of hosts generated for each simulator")
 	cmd.Flags().StringVarP(&ifaceName, "interface", "i", "", "network interface to use")
 	return cmd
 }
@@ -197,7 +203,7 @@ func run(simulators []simulatorInfo, extIP net.IP) error {
 		printMsg(s.name, "Starting")
 		printMsg(s.name, s.infoHeaders...)
 
-		hosts, err := s.s.Hosts()
+		hosts, err := s.s.Hosts(size)
 		if err != nil {
 			printMsg(s.name, color.RedString("failed: ")+err.Error())
 			continue

diff --git a/simulator/c2_dns.go b/simulator/c2_dns.go
@@ -33,7 +33,7 @@ func (*C2DNS) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns hosts marked c2 dns threat.
-func (t *C2DNS) Hosts() ([]string, error) {
+func (t *C2DNS) Hosts(size int) ([]string, error) {
 	resp, err := http.Get("https://cybercrime-tracker.net/all.php")
 	if err != nil {
 		return nil, errors.Wrapf(err, "cyber crime tracker get http")
@@ -54,7 +54,7 @@ func (t *C2DNS) Hosts() ([]string, error) {
 		return hosts, nil
 	}
 
-	for len(hosts) < 10 {
+	for len(hosts) < size {
 		c2URL := c2s[rand.Intn(len(c2s))]
 		u, err := url.Parse("http://" + c2URL)
 		if err != nil {

diff --git a/simulator/c2_ip.go b/simulator/c2_ip.go
@@ -34,8 +34,7 @@ func (*C2IP) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns hosts marked as c2 ip threat.
-func (t *C2IP) Hosts() ([]string, error) {
-	const nLookup = 10
+func (t *C2IP) Hosts(size int) ([]string, error) {
 	resp, err := http.Get("https://api.open.wisdom.alphasoc.net/v1/c2-ip")
 	if err != nil {
 		return nil, err
@@ -59,7 +58,7 @@ func (t *C2IP) Hosts() ([]string, error) {
 		hosts []string
 		idx   = rand.Perm(len(response.Hosts))
 	)
-	for n, i := 0, 0; n < len(response.Hosts) && i < nLookup; n, i = n+1, i+1 {
+	for n, i := 0, 0; n < len(response.Hosts) && i < size; n, i = n+1, i+1 {
 		hosts = append(hosts, response.Hosts[idx[n]])
 	}
 	return hosts, nil

diff --git a/simulator/dga.go b/simulator/dga.go
@@ -33,12 +33,11 @@ func (*DGA) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns random generated dga hosts.
-func (t *DGA) Hosts() ([]string, error) {
-	const nLookup = 5
+func (t *DGA) Hosts(size int) ([]string, error) {
 	var hosts []string
 
 	idx := rand.Perm(len(tlds))
-	for i := 0; i < nLookup; i++ {
+	for i := 0; i < size; i++ {
 		label := strings.ToLower(utils.RandString(7))
 		hosts = append(hosts, label+tlds[idx[0]])
 		hosts = append(hosts, label+tlds[idx[1]])

diff --git a/simulator/hijack.go b/simulator/hijack.go
@@ -37,6 +37,6 @@ func (*Hijack) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns one domain to simulate dns query.
-func (s *Hijack) Hosts() ([]string, error) {
+func (s *Hijack) Hosts(_ int) ([]string, error) {
 	return []string{"alphasoc.com"}, nil
 }
diff --git a/simulator/scan.go b/simulator/scan.go
@@ -27,14 +27,12 @@ var (
 
 // PortScan simulator.
 type PortScan struct {
-	hostNo int
 	portNo int
 }
 
 // NewPortScan creates port scan simulator.
 func NewPortScan() *PortScan {
 	return &PortScan{
-		hostNo: 10,
 		portNo: 10,
 	}
 }
@@ -54,13 +52,13 @@ func (*PortScan) Simulate(ctx context.Context, extIP net.IP, host string) error
 }
 
 // Hosts returns host:port generated from RFC 1918 addresses.
-func (s *PortScan) Hosts() ([]string, error) {
+func (s *PortScan) Hosts(size int) ([]string, error) {
 	var (
 		hosts []string
 		idx   = rand.Perm(len(scanPorts))
 	)
 
-	for i := 0; i < s.hostNo; i++ {
+	for i := 0; i < size; i++ {
 		ip := scanIPRanges[rand.Intn(len(scanIPRanges))]
 		ip.IP[len(ip.IP)-2] = byte(rand.Intn(256))
 		ip.IP[len(ip.IP)-1] = byte(rand.Intn(256))

diff --git a/simulator/simulator.go b/simulator/simulator.go
@@ -9,5 +9,5 @@ import (
 // traffic for different kind of threaths.
 type Simulator interface {
 	Simulate(ctx context.Context, extIP net.IP, host string) error
-	Hosts() ([]string, error)
+	Hosts(size int) ([]string, error)
 }
diff --git a/simulator/sink.go b/simulator/sink.go
@@ -34,8 +34,7 @@ func (*Sinkhole) Simulate(ctx context.Context, extIP net.IP, host string) error
 }
 
 // Hosts returns hosts marked as sinkhole threat.
-func (t *Sinkhole) Hosts() ([]string, error) {
-	const nLookup = 10
+func (t *Sinkhole) Hosts(size int) ([]string, error) {
 	resp, err := http.Get("https://api.open.wisdom.alphasoc.net/v1/sinkhole")
 	if err != nil {
 		return nil, err
@@ -59,7 +58,7 @@ func (t *Sinkhole) Hosts() ([]string, error) {
 		hosts []string
 		idx   = rand.Perm(len(response.Hosts))
 	)
-	for n, i := 0, 0; n < len(response.Hosts) && i < nLookup; n, i = n+1, i+1 {
+	for n, i := 0, 0; n < len(response.Hosts) && i < size; n, i = n+1, i+1 {
 		hosts = append(hosts, response.Hosts[idx[n]])
 	}
 	return hosts, nil

diff --git a/simulator/spambot.go b/simulator/spambot.go
@@ -79,14 +79,13 @@ func (*Spambot) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns host:port generated from RFC 1918 addresses.
-func (s *Spambot) Hosts() ([]string, error) {
-	const nLookup = 10
+func (s *Spambot) Hosts(size int) ([]string, error) {
 	var (
 		hosts []string
 		idx   = rand.Perm(len(domains))
 	)
 
-	for i, n := 0, 0; i < nLookup && n < len(domains); i, n = i+1, n+1 {
+	for i, n := 0, 0; i < size && n < len(domains); i, n = i+1, n+1 {
 		mx, err := net.LookupMX(domains[idx[n]])
 		if err != nil || len(mx) == 0 {
 			i--

diff --git a/simulator/tor.go b/simulator/tor.go
@@ -33,8 +33,7 @@ func (*Tor) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns tor exit nodes.
-func (s *Tor) Hosts() ([]string, error) {
-	const nLookup = 10
+func (s *Tor) Hosts(size int) ([]string, error) {
 
 	resp, err := http.Get("https://api.ipify.org")
 	if err != nil {
@@ -73,7 +72,7 @@ func (s *Tor) Hosts() ([]string, error) {
 		hosts []string
 		idx   = rand.Perm(len(ips))
 	)
-	for n, i := 0, 0; n < len(ips) && i < nLookup; n, i = n+1, i+1 {
+	for n, i := 0, 0; n < len(ips) && i < size; n, i = n+1, i+1 {
 		hosts = append(hosts, ips[idx[n]]+":80")
 	}
 

diff --git a/simulator/tunnel.go b/simulator/tunnel.go
@@ -30,11 +30,10 @@ func (*Tunnel) Simulate(ctx context.Context, extIP net.IP, host string) error {
 }
 
 // Hosts returns random generated hosts to alphasoc sandbox.
-func (t *Tunnel) Hosts() ([]string, error) {
-	const nLookup = 30
+func (t *Tunnel) Hosts(size int) ([]string, error) {
 	var hosts []string
 
-	for i := 0; i < nLookup; i++ {
+	for i := 0; i < size; i++ {
 		label := strings.ToLower(utils.RandString(30))
 		hosts = append(hosts, label+".sandbox.alphasoc.xyz")
 	}