Merge remote-tracking branch 'origin/add_yaml_support' into add_objectives_support

This commit is contained in:
John Donaldson 2019-08-13 22:01:07 +01:00
commit 4e19e39872
2 changed files with 96 additions and 52 deletions

View File

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

View File

@ -13,6 +13,7 @@ import random
import string import string
import tempfile import tempfile
import shlex import shlex
import yaml
messageChars = b'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' messageChars = b'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
@ -126,21 +127,64 @@ class Puzzle:
def read_stream(self, stream): def read_stream(self, stream):
header = True 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: for line in stream:
if header:
line = line.strip() line = line.strip()
if not line: if not line:
header = False break
continue
key, val = line.split(':', 1) key, val = line.split(':', 1)
key = key.lower() key = key.lower()
val = val.strip() val = val.strip()
self.handle_header_key(key, val)
def handle_header_key(self, key, val):
if key == 'author': if key == 'author':
self.authors.append(val) self.authors.append(val)
elif key == 'summary': elif key == 'summary':
self.summary = val self.summary = val
elif key == 'answer': elif key == 'answer':
if not isinstance(val, str):
raise ValueError("Answers must be strings, got %s, instead" % (type(val),))
self.answers.append(val) self.answers.append(val)
elif key == "answers":
for answer in val:
if not isinstance(answer, str):
raise ValueError("Answers must be strings, got %s, instead" % (type(answer),))
self.answers.append(answer)
elif key == 'pattern': elif key == 'pattern':
self.pattern = val self.pattern = val
elif key == 'hint': elif key == 'hint':
@ -177,8 +221,7 @@ class Puzzle:
self.ksas.append(val) self.ksas.append(val)
else: else:
raise ValueError("Unrecognized header field: {}".format(key)) raise ValueError("Unrecognized header field: {}".format(key))
else:
self.body.write(line)
def read_directory(self, path): def read_directory(self, path):
try: try: