From da294f7f1e656b0c2956313a5bc9b78d3b5cfa3a Mon Sep 17 00:00:00 2001
From: Justin Bastress <jbastress@censys.io>
Date: Mon, 29 Jan 2018 16:03:32 -0500
Subject: [PATCH 1/2] Add .templates folder in lieu of inline scripts; update
 templates to match more common scenarios

---
 integration_tests/.template/cleanup.sh |  9 ++++
 integration_tests/.template/schema.py  | 18 +++++++
 integration_tests/.template/setup.sh   | 27 +++++++++++
 integration_tests/.template/test.sh    | 23 +++++++++
 integration_tests/new.sh               | 66 +++++---------------------
 5 files changed, 88 insertions(+), 55 deletions(-)
 create mode 100755 integration_tests/.template/cleanup.sh
 create mode 100644 integration_tests/.template/schema.py
 create mode 100755 integration_tests/.template/setup.sh
 create mode 100755 integration_tests/.template/test.sh

diff --git a/integration_tests/.template/cleanup.sh b/integration_tests/.template/cleanup.sh
new file mode 100755
index 0000000..2eebf11
--- /dev/null
+++ b/integration_tests/.template/cleanup.sh
@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+set +e
+
+echo "#{MODULE_NAME}/cleanup: Tests cleanup for #{MODULE_NAME}"
+
+CONTAINER_NAME=zgrab_#{MODULE_NAME}
+
+docker stop $CONTAINER_NAME
diff --git a/integration_tests/.template/schema.py b/integration_tests/.template/schema.py
new file mode 100644
index 0000000..a2a56cc
--- /dev/null
+++ b/integration_tests/.template/schema.py
@@ -0,0 +1,18 @@
+# zschema sub-schema for zgrab2's #{MODULE_NAME} module
+# Registers zgrab2-#{MODULE_NAME} globally, and #{MODULE_NAME} with the main zgrab2 schema.
+from zschema.leaves import *
+from zschema.compounds import *
+import zschema.registry
+
+import schemas.zcrypto as zcrypto
+import schemas.zgrab2 as zgrab2
+
+#{MODULE_NAME}_scan_response = SubRecord({
+    "result": SubRecord({
+        # TODO FIXME IMPLEMENT SCHEMA
+    })
+}, extends = zgrab2.base_scan_response)
+
+zschema.registry.register_schema("zgrab2-#{MODULE_NAME}", #{MODULE_NAME}_scan_response)
+
+zgrab2.register_scan_response_type("#{MODULE_NAME}", #{MODULE_NAME}_scan_response)
diff --git a/integration_tests/.template/setup.sh b/integration_tests/.template/setup.sh
new file mode 100755
index 0000000..d0e8e7b
--- /dev/null
+++ b/integration_tests/.template/setup.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+
+echo "#{MODULE_NAME}/setup: Tests setup for #{MODULE_NAME}"
+
+# TODO FIXME -- set the container tag
+echo "TODO FIXME:  modify setup.sh to launch the integration test container"
+exit 1
+# CONTAINER_TAG=FIXME
+
+CONTAINER_NAME="zgrab_#{MODULE_NAME}"
+
+# If the container is already running, use it.
+if docker ps --filter "name=$CONTAINER_NAME" | grep -q $CONTAINER_NAME; then
+    echo "#{MODULE_NAME}/setup: Container $CONTAINER_NAME already running -- nothing to setup"
+    exit 0
+fi
+
+# If it is not running, try launching it -- on success, use that. 
+echo "#{MODULE_NAME}/setup: Trying to launch $CONTAINER_NAME..."
+if ! docker run --rm --name $CONTAINER_NAME -td $CONTAINER_TAG; then
+    echo "#{MODULE_NAME}/setup: Building docker image $CONTAINER_TAG..."
+    # If it fails, build it from ./container/Dockerfile
+    docker build -t $CONTAINER_TAG ./container
+    # Try again
+    echo "#{MODULE_NAME}/setup: Launching $CONTAINER_NAME..."
+    docker run --rm --name $CONTAINER_NAME -td $CONTAINER_TAG
+fi
diff --git a/integration_tests/.template/test.sh b/integration_tests/.template/test.sh
new file mode 100755
index 0000000..e0c8741
--- /dev/null
+++ b/integration_tests/.template/test.sh
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+set -e
+MODULE_DIR=$(dirname $0)
+ZGRAB_ROOT=$MODULE_DIR/../..
+ZGRAB_OUTPUT=$ZGRAB_ROOT/zgrab-output
+
+mkdir -p $ZGRAB_OUTPUT/#{MODULE_NAME}
+
+CONTAINER_NAME=zgrab_#{MODULE_NAME}
+
+OUTPUT_FILE=$ZGRAB_OUTPUT/#{MODULE_NAME}
+
+echo "#{MODULE_NAME}/test: Tests runner for #{MODULE_NAME}"
+# TODO FIXME: Add any necessary flags or additional tests
+CONTAINER_NAME=$CONTAINER_NAME $ZGRAB_ROOT/docker-runner/docker-run.sh #{MODULE_NAME} > $OUTPUT_FILE
+
+# Dump the docker logs
+echo "#{MODULE_NAME}/test: BEGIN docker logs from $CONTAINER_NAME [{("
+docker logs --tail all $CONTAINER_NAME
+echo ")}] END docker logs from $CONTAINER_NAME"
+
+# TODO: If there are any other relevant log files, dump those to stdout here.
diff --git a/integration_tests/new.sh b/integration_tests/new.sh
index 9e23992..92ed3f3 100755
--- a/integration_tests/new.sh
+++ b/integration_tests/new.sh
@@ -1,5 +1,6 @@
 #!/usr/bin/env bash
 
+set -e
 # Utility script for scaffolding stub test files for a new protocol
 
 # Run from root of project
@@ -16,62 +17,17 @@ module_name="$1"
 module_path="integration_tests/$module_name"
 
 mkdir -p $module_path
+pushd "integration_tests/.template"
+for file in $(ls *.sh); do
+    dest="../$module_name/$file"
+    cp "$file" "$dest"
+    sed -i "s/#{MODULE_NAME}/$module_name/g" "$dest"
+    chmod +x "$dest"
+done
+popd
 
-cat << EOF > $module_path/setup.sh
-#!/usr/bin/env bash
-
-echo "$module_name/setup: Tests setup for $module_name"
-EOF
-chmod +x $module_path/setup.sh
-
-cat << EOF > $module_path/test.sh
-#!/usr/bin/env bash
-
-set -e
-MODULE_DIR=\$(dirname \$0)
-TEST_ROOT=\$MODULE_DIR/..
-ZGRAB_ROOT=\$MODULE_DIR/../..
-ZGRAB_OUTPUT=\$ZGRAB_ROOT/zgrab-output
-
-mkdir -p \$ZGRAB_OUTPUT/$module_name
-
-# OUTPUT_FILE=[TODO].json
-
-echo "$module_name/test: Tests runner for $module_name"
-# CONTAINER_NAME=[TODO] \$ZGRAB_ROOT/docker-runner/docker-run.sh $module_name > \$OUTPUT_FILE
-
-EOF
-chmod +x $module_path/test.sh
-
-cat << EOF > $module_path/cleanup.sh
-#!/usr/bin/env bash
-
-set +e
-
-echo "$module_name/cleanup: Tests cleanup for $module_name"
-EOF
-chmod +x $module_path/cleanup.sh
-
-cat << EOF > schemas/$module_name.py
-# zschema sub-schema for zgrab2's $module_name module
-# Registers zgrab2-$module_name globally, and $module_name with the main zgrab2 schema.
-from zschema.leaves import *
-from zschema.compounds import *
-import zschema.registry
-
-import schemas.zcrypto as zcrypto
-import schemas.zgrab2 as zgrab2
-
-${module_name}_scan_response = SubRecord({
-    "result": SubRecord({
-        # TODO FIXME IMPLEMENT SCHEMA
-    })
-}, extends = zgrab2.base_scan_response)
-
-zschema.registry.register_schema("zgrab2-${module_name}", ${module_name}_scan_response)
-
-zgrab2.register_scan_response_type("${module_name}", ${module_name}_scan_response)
-EOF
+cp "integration_tests/.template/schema.py" "schemas/$module_name.py"
+sed -i "s/#{MODULE_NAME}/$module_name/g" "schemas/$module_name.py"
 
 echo "import schemas.$module_name" >> schemas/__init__.py
 

From 59aa3d857193c38cc8293bf92f5c7ff2c9233c19 Mon Sep 17 00:00:00 2001
From: Justin Bastress <jbastress@censys.io>
Date: Mon, 29 Jan 2018 16:36:45 -0500
Subject: [PATCH 2/2] Update README; fix output file in test

---
 integration_tests/.template/README.md | 29 +++++++++++++++++++++++++++
 integration_tests/.template/test.sh   |  2 +-
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 integration_tests/.template/README.md

diff --git a/integration_tests/.template/README.md b/integration_tests/.template/README.md
new file mode 100644
index 0000000..fa3d221
--- /dev/null
+++ b/integration_tests/.template/README.md
@@ -0,0 +1,29 @@
+Template integration test scripts, used by `new.sh`.
+
+ - All scripts start with the default `CONTAINER_NAME` variable matching
+   the standard format (`zgrab_<module name>`).
+
+ - `setup.sh` tries to launch the container; if it fails, then it tries
+   to build the docker image from the `Dockerfile` in
+   `integration_tests/<module_name>/container`; then it tries to launch
+   the container again. **This must be modified by all new modules**.
+
+ - `test.sh` runs the container using the `docker-runner.sh` script and
+   dumps the output in `zgrab-output/<module name>/<module_name>.json`.
+   Afterwards it dumps the container logs to stdout. This can be used
+   as-is, it only provides a single test case with no arguments.
+
+ - `cleanup.sh` runs `docker stop $CONTAINER_NAME`. This should tear
+   down everything started in `setup.sh`. In the case where there is
+   only a single container, it can be used as-is.
+
+ - `schema.py` registers the `scan_response_type` and provides a 
+   skeleton that can be filled out. **This must be modified for all new
+   modules**.
+
+These files have `#{MODULE_NAME}` statically replaced with the module's
+name (the argument passed to `new.sh`).
+
+ - `*.sh` get copied to `integration_tests/<module name>/*.sh`
+
+ - `schema.py` gets copied to `schemas/<module name>.py`
diff --git a/integration_tests/.template/test.sh b/integration_tests/.template/test.sh
index e0c8741..650cea1 100755
--- a/integration_tests/.template/test.sh
+++ b/integration_tests/.template/test.sh
@@ -9,7 +9,7 @@ mkdir -p $ZGRAB_OUTPUT/#{MODULE_NAME}
 
 CONTAINER_NAME=zgrab_#{MODULE_NAME}
 
-OUTPUT_FILE=$ZGRAB_OUTPUT/#{MODULE_NAME}
+OUTPUT_FILE=$ZGRAB_OUTPUT/#{MODULE_NAME}/#{MODULE_NAME}.json
 
 echo "#{MODULE_NAME}/test: Tests runner for #{MODULE_NAME}"
 # TODO FIXME: Add any necessary flags or additional tests