From a00cef992080f9cc364bf023adcac23c84b888ec Mon Sep 17 00:00:00 2001
From: Arne Welzel <arne.welzel@corelight.com>
Date: Fri, 21 Oct 2022 20:03:54 +0200
Subject: [PATCH] enum_names: Support naming types with a string

In his ZeekWeek 2022 presentation, @stevesmoot mentioned that he had a
difficult time looking up enum names when all he had was a string
naming the type.

Add support to enum_names() to transparently lookup the type if a string
is provided. This is similar in how record_fields() behaves when being
passed a string.
---
 NEWS                                       |  3 +++
 src/zeek.bif                               | 18 ++++++++++++------
 testing/btest/Baseline/bifs.enum_names/out | 20 ++++++++++++++++++++
 testing/btest/bifs/enum_names.zeek         | 12 ++++++++++++
 4 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/NEWS b/NEWS
index b31b1e62e3..c7c0b88007 100644
--- a/NEWS
+++ b/NEWS
@@ -46,6 +46,9 @@ Changed Functionality
   will be raised once only. Further, analyzer confirmations are not raised
   after a violation.
 
+- The parameter given to ``enum_names()`` can now be a string naming the
+  enum type, rather than the type itself.
+
 Deprecated Functionality
 ------------------------
 
diff --git a/src/zeek.bif b/src/zeek.bif
index d01c1c9b9e..c250a6ae18 100644
--- a/src/zeek.bif
+++ b/src/zeek.bif
@@ -1938,20 +1938,26 @@ function type_aliases%(x: any%): string_set
 
 ## Returns all value names associated with an enum type.
 ##
-## et: An enum type.
+## et: An enum type or a string naming one.
 ##
 ## Returns: All enum value names associated with enum type *et*.
-##          If *et* is not an enum type, an empty set is returned.
+##          If *et* is not an enum type or does not name one, an empty set is returned.
 function enum_names%(et: any%): string_set
 	%{
 	auto rval = make_intrusive<zeek::TableVal>(zeek::id::string_set);
 
-	if ( et->GetType()->Tag() != TYPE_TYPE )
-		return rval;
+	zeek::TypePtr t = zeek::Type::nil;
 
-	const auto& t = et->GetType()->AsTypeType()->GetType();
+	if ( et->GetType()->Tag() == TYPE_STRING )
+		{
+		const auto& id = zeek::detail::global_scope()->Find(et->AsStringVal()->ToStdString());
+		if ( id && id->IsType() )
+			t = id->GetType();
+		}
+	else if ( et->GetType()->Tag() == TYPE_TYPE )
+		t = et->GetType()->AsTypeType()->GetType();
 
-	if ( t->Tag() != TYPE_ENUM )
+	if ( ! t || t->Tag() != TYPE_ENUM )
 		return rval;
 
 	auto enum_type = t->AsEnumType();
diff --git a/testing/btest/Baseline/bifs.enum_names/out b/testing/btest/Baseline/bifs.enum_names/out
index 664204b545..1780b878fa 100644
--- a/testing/btest/Baseline/bifs.enum_names/out
+++ b/testing/btest/Baseline/bifs.enum_names/out
@@ -1,4 +1,20 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
+with types
+{
+Monochrome::WHITE,
+Monochrome::BLACK
+}
+{
+RED,
+GREEN,
+PURPLE,
+BLUE
+}
+{
+Monochrome::WHITE,
+Monochrome::BLACK
+}
+with strings
 {
 RED,
 GREEN,
@@ -11,3 +27,7 @@ GREEN,
 PURPLE,
 BLUE
 }
+{
+Monochrome::WHITE,
+Monochrome::BLACK
+}
diff --git a/testing/btest/bifs/enum_names.zeek b/testing/btest/bifs/enum_names.zeek
index 52ef537664..0cb8bac0c5 100644
--- a/testing/btest/bifs/enum_names.zeek
+++ b/testing/btest/bifs/enum_names.zeek
@@ -6,5 +6,17 @@ type ColorAlias: Color;
 
 redef enum Color += { PURPLE };
 
+module Monochrome;
+
+type Color: enum { WHITE, BLACK };
+
+
+print "with types";
 print enum_names(Color);
 print enum_names(ColorAlias);
+print enum_names(Monochrome::Color);
+
+print "with strings";
+print enum_names("Color");
+print enum_names("ColorAlias");
+print enum_names("Monochrome::Color");