From 6b285e611cfc7f5c6db6efe4ee9f70f5bb784ae1 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Tue, 20 Mar 2018 09:49:03 +0200
Subject: [PATCH] dvc: validate stage files

Fixes #578

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/output.py    |  9 +++++++++
 dvc/stage.py     | 23 +++++++++++++++++++++++
 requirements.txt |  1 +
 setup.py         |  1 +
 4 files changed, 34 insertions(+)

diff --git a/dvc/output.py b/dvc/output.py
index 39a0d08a5..a17d469c3 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -2,6 +2,7 @@ import os
 import stat
 import json
 import shutil
+import schema
 from checksumdir import dirhash
 
 from dvc.system import System
@@ -50,6 +51,11 @@ class Dependency(object):
     PARAM_MD5 = 'md5'
     MD5_DIR_SUFFIX = '.dir'
 
+    SCHEMA = {
+        PARAM_PATH: str,
+        schema.Optional(PARAM_MD5): schema.Or(str, None),
+    }
+
     def __init__(self, project, path, md5=None):
         self.project = project
         self.path = os.path.abspath(os.path.realpath(path))
@@ -126,6 +132,9 @@ class Dependency(object):
 class Output(Dependency):
     PARAM_CACHE = 'cache'
 
+    SCHEMA = Dependency.SCHEMA
+    SCHEMA[schema.Optional(PARAM_CACHE)] = bool
+
     def __init__(self, project, path, md5=None, use_cache=True):
         super(Output, self).__init__(project, path, md5=md5)
         self.use_cache = use_cache
diff --git a/dvc/stage.py b/dvc/stage.py
index 9c9af07be..3a369ca37 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -2,9 +2,11 @@ import os
 import yaml
 import itertools
 import subprocess
+import schema
 
 from dvc.exceptions import DvcException
 from dvc.output import Output, Dependency, OutputError
+from dvc.logger import Logger
 
 
 class StageCmdFailedError(DvcException):
@@ -13,6 +15,11 @@ class StageCmdFailedError(DvcException):
         super(StageCmdFailedError, self).__init__(msg)
 
 
+class StageFileFormatError(DvcException):
+    def __init__(self):
+        super(StageFileFormatError, self).__init__('Stage file format error')
+
+
 class MissingDataSource(OutputError):
     def __init__(self, missing_files):
         assert len(missing_files) > 0
@@ -33,6 +40,12 @@ class Stage(object):
     PARAM_DEPS = 'deps'
     PARAM_OUTS = 'outs'
 
+    SCHEMA = {
+        schema.Optional(PARAM_CMD): schema.Or(str, None),
+        schema.Optional(PARAM_DEPS): schema.Or(schema.And(list, schema.Schema([Dependency.SCHEMA])), None),
+        schema.Optional(PARAM_OUTS): schema.Or(schema.And(list, schema.Schema([Output.SCHEMA])), None),
+    }
+
     def __init__(self, project, path=None, cmd=None, cwd=None, deps=[], outs=[]):
         self.project = project
         self.path = path
@@ -94,8 +107,18 @@ class Stage(object):
 
         return self
 
+    @staticmethod
+    def validate(d):
+        try:
+            schema.Schema(Stage.SCHEMA).validate(d)
+        except schema.SchemaError as exc:
+            Logger.debug(str(exc))
+            raise StageFileFormatError()
+
     @staticmethod
     def loadd(project, d, path):
+        Stage.validate(d)
+
         path = os.path.abspath(path)
         cwd = os.path.dirname(path)
         cmd = d.get(Stage.PARAM_CMD, None)
diff --git a/requirements.txt b/requirements.txt
index 77574779a..3ac6c3e0b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,3 +15,4 @@ checksumdir>=1.1.4
 setuptools>=34.0.0
 nanotime>=0.5.2
 pyasn1>=0.4.1
+schema>=0.6.7
diff --git a/setup.py b/setup.py
index 153297cd5..f8f6b7018 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,7 @@ install_requires = [
     "setuptools>=34.0.0",
     "nanotime>=0.5.2",
     "pyasn1>=0.4.1",
+    "schema>=0.6.7",
 ]
 
 setup(