From cc2b7983759aa0074afbac4254f11c961455c27e Mon Sep 17 00:00:00 2001
From: dmiller <dmiller@e0a8ed71-7df4-0310-8962-fdc924857419>
Date: Wed, 1 May 2024 22:19:31 +0000
Subject: [PATCH] Make NetBlockRandomIPv4 responsible for tracking number of
 random addrs

---
 TargetGroup.cc | 13 +++++++++++--
 TargetGroup.h  |  6 +++++-
 targets.cc     | 12 ++++--------
 targets.h      |  3 +--
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/TargetGroup.cc b/TargetGroup.cc
index d66b198a6..21fe58d55 100644
--- a/TargetGroup.cc
+++ b/TargetGroup.cc
@@ -111,6 +111,7 @@ public:
    * the return value to the pointer that this method was called through.
    * On error, return NULL. */
   virtual NetBlock *resolve() { return this; }
+  virtual void reject_last_host() {}
   virtual bool next(struct sockaddr_storage *ss, size_t *sslen) = 0;
   virtual void apply_netmask(int bits) = 0;
   virtual std::string str() const = 0;
@@ -120,12 +121,15 @@ class NetBlockRandomIPv4 : public NetBlock {
 public:
   NetBlockRandomIPv4();
 
+  void reject_last_host() { count++; }
+  void set_num_random(int num) { count = num; }
   bool next(struct sockaddr_storage *ss, size_t *sslen);
   void apply_netmask(int bits) {}
   std::string str() const {return "Random IPv4 addresses";}
 
 private:
   struct sockaddr_in base;
+  int count;
 };
 
 class NetBlockIPv4Ranges : public NetBlock {
@@ -337,17 +341,21 @@ bool NetBlock::is_resolved_address(const struct sockaddr_storage *ss) const {
   return false;
 }
 
-NetBlockRandomIPv4::NetBlockRandomIPv4() {
+NetBlockRandomIPv4::NetBlockRandomIPv4() : count(0) {
   memset(&base, 0, sizeof(base));
   base.sin_family = AF_INET;
 }
 
 bool NetBlockRandomIPv4::next(struct sockaddr_storage *ss, size_t *sslen) {
+  if (count <= 0) {
+    return false;
+  }
   do {
     base.sin_addr.s_addr = get_random_unique_u32();
   } while (ip_is_reserved(&base.sin_addr));
   memcpy(ss, &base, sizeof(base));
   *sslen = sizeof(base);
+  count--;
   return true;
 }
 
@@ -803,9 +811,10 @@ int TargetGroup::parse_expr(const char *target_expr, int af) {
     return 1;
 }
 
-void TargetGroup::generate_random_ips() {
+void TargetGroup::generate_random_ips(int num_random) {
   assert(this->netblock == NULL);
   this->netblock = new NetBlockRandomIPv4();
+  this->netblock->set_num_random(num_random);
 }
 
 /* Grab the next host from this expression (if any) and updates its internal
diff --git a/TargetGroup.h b/TargetGroup.h
index 2430bba61..61daae8c4 100644
--- a/TargetGroup.h
+++ b/TargetGroup.h
@@ -101,7 +101,11 @@ public:
   const std::list<struct sockaddr_storage> &get_unscanned_addrs(void) const;
   /* is the current expression a named host */
   int get_namedhost() const;
-  void generate_random_ips();
+  void generate_random_ips(int num_random);
+  void reject_last_host();
+
+  private:
+  std::list<NetBlock *>netblocks;
 };
 
 #endif /* TARGETGROUP_H */
diff --git a/targets.cc b/targets.cc
index 3a4b33892..60950d7fa 100644
--- a/targets.cc
+++ b/targets.cc
@@ -296,9 +296,8 @@ HostGroupState::HostGroupState(int lookahead, int rnd, int num_random, int argc,
   current_batch_sz = 0;
   next_batch_no = 0;
   randomize = rnd;
-  this->num_random = num_random;
   if (num_random > 0) {
-    current_group.generate_random_ips();
+    current_group.generate_random_ips(num_random);
   }
 }
 
@@ -444,15 +443,12 @@ bool HostGroupState::get_next_host(struct sockaddr_storage *ss, size_t *sslen, s
       }
     }
     /* Check exclude list. */
-    if (!hostInExclude((struct sockaddr *) ss, *sslen, exclude_group))
+    if (!hostInExclude((struct sockaddr *) ss, *sslen, exclude_group)) {
+      current_group.reject_last_host();
       break;
+    }
   } while (true);
 
-  // If this is the last of the random IPs, pop in a new expression
-  if (num_random > 0 && --num_random == 0) {
-    process_next_expression();
-  }
-
   return true;
 }
 
diff --git a/targets.h b/targets.h
index a6dda3ff2..c6e10a5aa 100644
--- a/targets.h
+++ b/targets.h
@@ -100,11 +100,10 @@ public:
   /* Returns true iff the defer buffer is not yet full. */
   bool defer(Target *t);
   void undefer();
+  const char *next_expression();
   bool get_next_host(struct sockaddr_storage *ss, size_t *sslen, struct addrset *exclude_group);
   private:
-  int num_random;
   bool process_next_expression();
-  const char *next_expression();
 };
 
 /* ports is used to pass information about what ports to use for host discovery */