From f147c472714b68654a1f2699c1b28e39f6fd67fd Mon Sep 17 00:00:00 2001
From: Jon Siwek <jsiwek@corelight.com>
Date: Fri, 6 Nov 2020 18:31:41 -0800
Subject: [PATCH] Add enum_names() BIF to return names of an enum type's values

---
 src/zeek.bif                               | 26 ++++++++++++++++++++++
 testing/btest/Baseline/bifs.enum_names/out | 13 +++++++++++
 testing/btest/bifs/enum_names.zeek         | 10 +++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 testing/btest/Baseline/bifs.enum_names/out
 create mode 100644 testing/btest/bifs/enum_names.zeek

diff --git a/src/zeek.bif b/src/zeek.bif
index b9febd0df0..b55948efbe 100644
--- a/src/zeek.bif
+++ b/src/zeek.bif
@@ -1946,6 +1946,32 @@ function type_aliases%(x: any%): string_set
 	return rval;
 	%}
 
+## Returns all value names associated with an enum type.
+##
+## et: An enum type.
+##
+## Returns: All enum value names associated with enum type *et*.
+##          If *et* is not an enum type, 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;
+
+	const auto& t = et->GetType()->AsTypeType()->GetType();
+
+	if ( t->Tag() != TYPE_ENUM )
+		return rval;
+
+	auto enum_type = t->AsEnumType();
+
+	for ( const auto& [name, i] : enum_type->Names() )
+		rval->Assign(make_intrusive<zeek::StringVal>(name), nullptr);
+
+	return rval;
+	%}
+
 ## Returns: list of command-line arguments (``argv``) used to run Zeek.
 function zeek_args%(%): string_vec
 	%{
diff --git a/testing/btest/Baseline/bifs.enum_names/out b/testing/btest/Baseline/bifs.enum_names/out
new file mode 100644
index 0000000000..664204b545
--- /dev/null
+++ b/testing/btest/Baseline/bifs.enum_names/out
@@ -0,0 +1,13 @@
+### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
+{
+RED,
+GREEN,
+PURPLE,
+BLUE
+}
+{
+RED,
+GREEN,
+PURPLE,
+BLUE
+}
diff --git a/testing/btest/bifs/enum_names.zeek b/testing/btest/bifs/enum_names.zeek
new file mode 100644
index 0000000000..52ef537664
--- /dev/null
+++ b/testing/btest/bifs/enum_names.zeek
@@ -0,0 +1,10 @@
+# @TEST-EXEC: zeek -b %INPUT >out
+# @TEST-EXEC: btest-diff out
+
+type Color: enum { RED, GREEN, BLUE };
+type ColorAlias: Color;
+
+redef enum Color += { PURPLE };
+
+print enum_names(Color);
+print enum_names(ColorAlias);