diff --git a/libpathod/language.py b/libpathod/language.py index 1445ca82b..a5c53677a 100644 --- a/libpathod/language.py +++ b/libpathod/language.py @@ -1,4 +1,4 @@ -import operator, string, random, mmap, os, time +import operator, string, random, mmap, os, time, copy from email.utils import formatdate import contrib.pyparsing as pp from netlib import http_status, tcp @@ -189,7 +189,7 @@ class _Value: return LiteralGenerator(self.val) def __repr__(self): - return self.val + return self.spec() class ValueLiteral(_Value): @@ -198,6 +198,9 @@ class ValueLiteral(_Value): e = v_literal.copy() return e.setParseAction(lambda x: klass(*x)) + def spec(self): + return '"%s"'%self.val.encode("string_escape") + class ValueNakedLiteral(_Value): @classmethod @@ -205,6 +208,9 @@ class ValueNakedLiteral(_Value): e = v_naked_literal.copy() return e.setParseAction(lambda x: klass(*x)) + def spec(self): + return self.val.encode("string_escape") + class ValueGenerate: def __init__(self, usize, unit, datatype): @@ -230,8 +236,16 @@ class ValueGenerate: e += pp.Optional(s, default="bytes") return e.setParseAction(lambda x: klass(*x)) - def __str__(self): - return "@%s%s,%s"%(self.usize, self.unit, self.datatype) + def spec(self): + s = "@%s"%self.usize + if self.unit != "b": + s += self.unit + if self.datatype != "bytes": + s += ",%s"%self.datatype + return s + + def __repr__(self): + return self.spec() class ValueFile: @@ -259,8 +273,8 @@ class ValueFile: raise FileAccessDenied("File not readable") return FileGenerator(s) - def __str__(self): - return "<%s"%(self.path) + def spec(self): + return '<"%s"'%self.path.encode("string_escape") Value = pp.MatchFirst( @@ -410,9 +424,24 @@ class _Action: def __init__(self, offset): self.offset = offset + def resolve_offset(self, msg): + """ + Resolves offset specifications to a numeric offset. Returns a copy + of the action object. + """ + c = copy.copy(self) + if c.offset == "r": + c.offset = random.randrange(msg.length()) + elif c.offset == "a": + c.offset = msg.length() + 1 + return c + def __cmp__(self, other): return cmp(self.offset, other.offset) + def __repr__(self): + return self.spec() + class PauseAt(_Action): def __init__(self, offset, seconds): @@ -432,6 +461,9 @@ class PauseAt(_Action): ) return e.setParseAction(lambda x: klass(*x)) + def spec(self): + return "p%s,%s"%(self.offset, self.seconds) + def accept(self, settings, r): r.actions.append((self.offset, "pause", self.seconds)) @@ -449,6 +481,9 @@ class DisconnectAt(_Action): e += Offset return e.setParseAction(lambda x: klass(*x)) + def spec(self): + return "d%s"%self.offset + class InjectAt(_Action): def __init__(self, offset, value): @@ -463,6 +498,9 @@ class InjectAt(_Action): e += Value return e.setParseAction(lambda x: klass(*x)) + def spec(self): + return "i%s,%s"%(self.offset, self.value.spec()) + def accept(self, settings, r): r.actions.append( ( diff --git a/test/test_language.py b/test/test_language.py index 7680492f7..e1697f954 100644 --- a/test/test_language.py +++ b/test/test_language.py @@ -5,6 +5,106 @@ import tutils language.TESTING = True +class TestValueNakedLiteral: + def test_expr(self): + v = language.ValueNakedLiteral("foo") + assert v.expr() + + def test_spec(self): + v = language.ValueNakedLiteral("foo") + assert v.spec() == repr(v) == "foo" + + v = language.ValueNakedLiteral("f\x00oo") + assert v.spec() == repr(v) == r"f\x00oo" + + +class TestValueLiteral: + def test_espr(self): + v = language.ValueLiteral("foo") + assert v.expr() + assert v.val == "foo" + + v = language.ValueLiteral(r"foo\n") + assert v.expr() + assert v.val == "foo\n" + assert repr(v) + + def test_spec(self): + v = language.ValueLiteral("foo") + assert v.spec() == r'"foo"' + + v = language.ValueLiteral("f\x00oo") + assert v.spec() == repr(v) == r'"f\x00oo"' + + +class TestValueGenerate: + def test_basic(self): + v = language.Value.parseString("@10b")[0] + assert v.usize == 10 + assert v.unit == "b" + assert v.bytes() == 10 + v = language.Value.parseString("@10")[0] + assert v.unit == "b" + v = language.Value.parseString("@10k")[0] + assert v.bytes() == 10240 + v = language.Value.parseString("@10g")[0] + assert v.bytes() == 1024**3 * 10 + + v = language.Value.parseString("@10g,digits")[0] + assert v.datatype == "digits" + g = v.get_generator({}) + assert g[:100] + + v = language.Value.parseString("@10,digits")[0] + assert v.unit == "b" + assert v.datatype == "digits" + + def test_spec(self): + v = language.ValueGenerate(1, "b", "bytes") + assert v.spec() == repr(v) == "@1" + + v = language.ValueGenerate(1, "k", "bytes") + assert v.spec() == repr(v) == "@1k" + + v = language.ValueGenerate(1, "k", "ascii") + assert v.spec() == repr(v) == "@1k,ascii" + + v = language.ValueGenerate(1, "b", "ascii") + assert v.spec() == repr(v) == "@1,ascii" + + +class TestValueFile: + def test_file_value(self): + v = language.Value.parseString("<'one two'")[0] + assert str(v) + assert v.path == "one two" + + v = language.Value.parseString("