Adding YAML support to Moth files

Dockerfile.moth-devel

@@ -9,7 +9,8 @@ RUN apk --no-cache add \
     && \
     pip3 install \
       scapy==2.4.2 \
-      pillow==5.4.1
+      pillow==5.4.1 \
+      PyYAML==5.1.1
 
 
 COPY devel /app/

devel/moth.py

@@ -13,6 +13,7 @@ import random
 import string
 import tempfile
 import shlex
+import yaml
 
 messageChars = b'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
@@ -91,21 +92,62 @@ class Puzzle:
 
     def read_stream(self, stream):
         header = True
+        line = ""
+        if stream.read(3) == "---":
+            header = "yaml"
+        else:
+            header = "moth"
+
+        stream.seek(0)
+
+        if header == "yaml":
+            self.read_yaml_header(stream)
+        elif header == "moth":
+            self.read_moth_header(stream)
+                
+        for line in stream:
+            self.body.write(line)
+
+    def read_yaml_header(self, stream):
+        contents = ""
+        header = False
+        for line in stream:
+            if line.strip() == "---" and header:  # Handle last line
+                break
+            elif line.strip() == "---":  # Handle first line
+                header = True
+                continue
+            else:
+                contents += line
+
+        config = yaml.safe_load(contents)
+        for key, value in config.items():
+            key = key.lower()
+            self.handle_header_key(key, value)
+
+
+    def read_moth_header(self, stream):
         for line in stream:
-            if header:
             line = line.strip()
             if not line:
-                    header = False
+                break
-                    continue
+
             key, val = line.split(':', 1)
             key = key.lower()
             val = val.strip()
+            self.handle_header_key(key, val)
+
+    def handle_header_key(self, key, val):
         if key == 'author':
             self.authors.append(val)
         elif key == 'summary':
             self.summary = val
         elif key == 'answer':
             self.answers.append(val)
+        elif key == "answers":
+            for answer in val:
+                answer = str(answer)
+                self.answers.append(answer)
         elif key == 'pattern':
             self.pattern = val
         elif key == 'hint':
@@ -130,8 +172,7 @@ class Puzzle:
             self.scripts.append(val)
         else:
             raise ValueError("Unrecognized header field: {}".format(key))
-            else:
+
-                self.body.write(line)
 
     def read_directory(self, path):
         try: