1 files changed, 16 insertions, 4 deletions
diff --git a/internal/behavior/store.go b/internal/behavior/store.go
index d31fe84..3ea527c 100644
--- a/internal/behavior/store.go
+++ b/internal/behavior/store.go
@@ -32,13 +32,16 @@ func excludedContentPath(root, path string) bool {
 func BuildPathIndex(dir string) map[string]string {
 	index := make(map[string]string)
 	filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
-		if err != nil || d.IsDir() || !strings.HasSuffix(d.Name(), ".yaml") {
+		if err != nil || d.IsDir() {
+			return nil
+		}
+		id, ok := yamlID(d.Name())
+		if !ok {
 			return nil
 		}
 		if excludedContentPath(dir, path) {
 			return nil
 		}
-		id := strings.TrimSuffix(d.Name(), ".yaml")
 		index[id] = path
 		return nil
 	})
@@ -48,13 +51,16 @@ func BuildPathIndex(dir string) map[string]string {
 func BuildRoomIndex(dir string) map[int]string {
 	index := make(map[int]string)
 	filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
-		if err != nil || d.IsDir() || !strings.HasSuffix(d.Name(), ".yaml") {
+		if err != nil || d.IsDir() {
+			return nil
+		}
+		id, ok := yamlID(d.Name())
+		if !ok {
 			return nil
 		}
 		if excludedContentPath(dir, path) {
 			return nil
 		}
-		id := strings.TrimSuffix(d.Name(), ".yaml")
 		n, convErr := strconv.Atoi(id)
 		if convErr != nil {
 			// A non-numeric filename (for example a "1001_copy.yaml"
@@ -187,6 +193,9 @@ func LoadDropTable(dataDir, id string) (*DropTableDef, error) {
 	index := loadDropIndex(dataDir)
 	path, ok := index[id]
 	if !ok {
+		if id == "" || strings.Contains(id, "..") || strings.ContainsAny(id, `/\`) {
+			return nil, fmt.Errorf("invalid drop table id %q", id)
+		}
 		path = filepath.Join(dataDir, "drops", id+".yaml")
 	}
 	data, err := os.ReadFile(path)
@@ -268,6 +277,9 @@ func resolveDrop(dataDir string, drops []DropEntry, depth int) *DropEntry {
 			Quantity:       qty,
 			Depletes:       picked.Depletes,
 			SuccessMessage: picked.SuccessMessage,
+			Tool:           picked.Tool,
+			Level:          picked.Level,
+			XP:             picked.XP,
 		}
 	}
 	return picked