util/stanza.lua

   1 -- Prosody IM
   2 -- Copyright (C) 2008-2010 Matthew Wild
   3 -- Copyright (C) 2008-2010 Waqas Hussain
   4 --
   5 -- This project is MIT/X11 licensed. Please see the
   6 -- COPYING file in the source package for more information.
   7 --
   8
   9
  10 local t_insert      =  table.insert;
  11 local t_concat      =  table.concat;
  12 local t_remove      =  table.remove;
  13 local t_concat      =  table.concat;
  14 local s_format      = string.format;
  15 local s_match       =  string.match;
  16 local tostring      =      tostring;
  17 local setmetatable  =  setmetatable;
  18 local getmetatable  =  getmetatable;
  19 local pairs         =         pairs;
  20 local ipairs        =        ipairs;
  21 local type          =          type;
  22 local next          =          next;
  23 local print         =         print;
  24 local unpack        =        unpack;
  25 local s_gsub        =   string.gsub;
  26 local s_char        =   string.char;
  27 local s_find        =   string.find;
  28 local os            =            os;
  29
  30 local do_pretty_printing = not os.getenv("WINDIR");
  31 local getstyle, getstring;
  32 if do_pretty_printing then
  33         local ok, termcolours = pcall(require, "util.termcolours");
  34         if ok then
  35                 getstyle, getstring = termcolours.getstyle, termcolours.getstring;
  36         else
  37                 do_pretty_printing = nil;
  38         end
  39 end
  40
  41 local xmlns_stanzas = "urn:ietf:params:xml:ns:xmpp-stanzas";
  42
  43 module "stanza"
  44
  45 stanza_mt = { __type = "stanza" };
  46 stanza_mt.__index = stanza_mt;
  47 local stanza_mt = stanza_mt;
  48
  49 function stanza(name, attr)
  50         local stanza = { name = name, attr = attr or {}, tags = {} };
  51         return setmetatable(stanza, stanza_mt);
  52 end
  53 local stanza = stanza;
  54
  55 function stanza_mt:query(xmlns)
  56         return self:tag("query", { xmlns = xmlns });
  57 end
  58
  59 function stanza_mt:body(text, attr)
  60         return self:tag("body", attr):text(text);
  61 end
  62
  63 function stanza_mt:tag(name, attrs)
  64         local s = stanza(name, attrs);
  65         local last_add = self.last_add;
  66         if not last_add then last_add = {}; self.last_add = last_add; end
  67         (last_add[#last_add] or self):add_direct_child(s);
  68         t_insert(last_add, s);
  69         return self;
  70 end
  71
  72 function stanza_mt:text(text)
  73         local last_add = self.last_add;
  74         (last_add and last_add[#last_add] or self):add_direct_child(text);
  75         return self;
  76 end
  77
  78 function stanza_mt:up()
  79         local last_add = self.last_add;
  80         if last_add then t_remove(last_add); end
  81         return self;
  82 end
  83
  84 function stanza_mt:reset()
  85         self.last_add = nil;
  86         return self;
  87 end
  88
  89 function stanza_mt:add_direct_child(child)
  90         if type(child) == "table" then
  91                 t_insert(self.tags, child);
  92         end
  93         t_insert(self, child);
  94 end
  95
  96 function stanza_mt:add_child(child)
  97         local last_add = self.last_add;
  98         (last_add and last_add[#last_add] or self):add_direct_child(child);
  99         return self;
 100 end
 101
 102 function stanza_mt:get_child(name, xmlns)
 103         for _, child in ipairs(self.tags) do
 104                 if (not name or child.name == name)
 105                         and ((not xmlns and self.attr.xmlns == child.attr.xmlns)
 106                                 or child.attr.xmlns == xmlns) then
 107
 108                         return child;
 109                 end
 110         end
 111 end
 112
 113 function stanza_mt:get_child_text(name, xmlns)
 114         local tag = self:get_child(name, xmlns);
 115         if tag then
 116                 return tag:get_text();
 117         end
 118         return nil;
 119 end
 120
 121 function stanza_mt:child_with_name(name)
 122         for _, child in ipairs(self.tags) do
 123                 if child.name == name then return child; end
 124         end
 125 end
 126
 127 function stanza_mt:child_with_ns(ns)
 128         for _, child in ipairs(self.tags) do
 129                 if child.attr.xmlns == ns then return child; end
 130         end
 131 end
 132
 133 function stanza_mt:children()
 134         local i = 0;
 135         return function (a)
 136                         i = i + 1
 137                         return a[i];
 138                 end, self, i;
 139 end
 140
 141 function stanza_mt:matching_tags(name, xmlns)
 142         xmlns = xmlns or self.attr.xmlns;
 143         local tags = self.tags;
 144         local start_i, max_i = 1, #tags;
 145         return function ()
 146                         for i=start_i,max_i do
 147                                 v = tags[i];
 148                                 if (not name or v.name == name)
 149                                 and (not xmlns or xmlns == v.attr.xmlns) then
 150                                         start_i = i+1;
 151                                         return v;
 152                                 end
 153                         end
 154                 end, tags, i;
 155 end
 156
 157 function stanza_mt:childtags()
 158         local i = 0;
 159         return function (a)
 160                         i = i + 1
 161                         local v = self.tags[i]
 162                         if v then return v; end
 163                 end, self.tags[1], i;
 164 end
 165
 166 function stanza_mt:maptags(callback)
 167         local tags, curr_tag = self.tags, 1;
 168         local n_children, n_tags = #self, #tags;
 169
 170         local i = 1;
 171         while curr_tag <= n_tags do
 172                 if self[i] == tags[curr_tag] then
 173                         local ret = callback(self[i]);
 174                         if ret == nil then
 175                                 t_remove(self, i);
 176                                 t_remove(tags, curr_tag);
 177                                 n_children = n_children - 1;
 178                                 n_tags = n_tags - 1;
 179                         else
 180                                 self[i] = ret;
 181                                 tags[i] = ret;
 182                         end
 183                         i = i + 1;
 184                         curr_tag = curr_tag + 1;
 185                 end
 186         end
 187         return self;
 188 end
 189
 190 local xml_escape
 191 do
 192         local escape_table = { ["'"] = "&apos;", ["\""] = "&quot;", ["<"] = "&lt;", [">"] = "&gt;", ["&"] = "&amp;" };
 193         function xml_escape(str) return (s_gsub(str, "['&<>\"]", escape_table)); end
 194         _M.xml_escape = xml_escape;
 195 end
 196
 197 local function _dostring(t, buf, self, xml_escape, parentns)
 198         local nsid = 0;
 199         local name = t.name
 200         t_insert(buf, "<"..name);
 201         for k, v in pairs(t.attr) do
 202                 if s_find(k, "\1", 1, true) then
 203                         local ns, attrk = s_match(k, "^([^\1]*)\1?(.*)$");
 204                         nsid = nsid + 1;
 205                         t_insert(buf, " xmlns:ns"..nsid.."='"..xml_escape(ns).."' ".."ns"..nsid..":"..attrk.."='"..xml_escape(v).."'");
 206                 elseif not(k == "xmlns" and v == parentns) then
 207                         t_insert(buf, " "..k.."='"..xml_escape(v).."'");
 208                 end
 209         end
 210         local len = #t;
 211         if len == 0 then
 212                 t_insert(buf, "/>");
 213         else
 214                 t_insert(buf, ">");
 215                 for n=1,len do
 216                         local child = t[n];
 217                         if child.name then
 218                                 self(child, buf, self, xml_escape, t.attr.xmlns);
 219                         else
 220                                 t_insert(buf, xml_escape(child));
 221                         end
 222                 end
 223                 t_insert(buf, "</"..name..">");
 224         end
 225 end
 226 function stanza_mt.__tostring(t)
 227         local buf = {};
 228         _dostring(t, buf, _dostring, xml_escape, nil);
 229         return t_concat(buf);
 230 end
 231
 232 function stanza_mt.top_tag(t)
 233         local attr_string = "";
 234         if t.attr then
 235                 for k, v in pairs(t.attr) do if type(k) == "string" then attr_string = attr_string .. s_format(" %s='%s'", k, xml_escape(tostring(v))); end end
 236         end
 237         return s_format("<%s%s>", t.name, attr_string);
 238 end
 239
 240 function stanza_mt.get_text(t)
 241         if #t.tags == 0 then
 242                 return t_concat(t);
 243         end
 244 end
 245
 246 function stanza_mt.get_error(stanza)
 247         local type, condition, text;
 248
 249         local error_tag = stanza:get_child("error");
 250         if not error_tag then
 251                 return nil, nil, nil;
 252         end
 253         type = error_tag.attr.type;
 254
 255         for child in error_tag:childtags() do
 256                 if child.attr.xmlns == xmlns_stanzas then
 257                         if not text and child.name == "text" then
 258                                 text = child:get_text();
 259                         elseif not condition then
 260                                 condition = child.name;
 261                         end
 262                         if condition and text then
 263                                 break;
 264                         end
 265                 end
 266         end
 267         return type, condition or "undefined-condition", text;
 268 end
 269
 270 function stanza_mt.__add(s1, s2)
 271         return s1:add_direct_child(s2);
 272 end
 273
 274
 275 do
 276         local id = 0;
 277         function new_id()
 278                 id = id + 1;
 279                 return "lx"..id;
 280         end
 281 end
 282
 283 function preserialize(stanza)
 284         local s = { name = stanza.name, attr = stanza.attr };
 285         for _, child in ipairs(stanza) do
 286                 if type(child) == "table" then
 287                         t_insert(s, preserialize(child));
 288                 else
 289                         t_insert(s, child);
 290                 end
 291         end
 292         return s;
 293 end
 294
 295 function deserialize(stanza)
 296         -- Set metatable
 297         if stanza then
 298                 local attr = stanza.attr;
 299                 for i=1,#attr do attr[i] = nil; end
 300                 local attrx = {};
 301                 for att in pairs(attr) do
 302                         if s_find(att, "|", 1, true) and not s_find(att, "\1", 1, true) then
 303                                 local ns,na = s_match(att, "^([^|]+)|(.+)$");
 304                                 attrx[ns.."\1"..na] = attr[att];
 305                                 attr[att] = nil;
 306                         end
 307                 end
 308                 for a,v in pairs(attrx) do
 309                         attr[a] = v;
 310                 end
 311                 setmetatable(stanza, stanza_mt);
 312                 for _, child in ipairs(stanza) do
 313                         if type(child) == "table" then
 314                                 deserialize(child);
 315                         end
 316                 end
 317                 if not stanza.tags then
 318                         -- Rebuild tags
 319                         local tags = {};
 320                         for _, child in ipairs(stanza) do
 321                                 if type(child) == "table" then
 322                                         t_insert(tags, child);
 323                                 end
 324                         end
 325                         stanza.tags = tags;
 326                 end
 327         end
 328
 329         return stanza;
 330 end
 331
 332 local function _clone(stanza)
 333         local attr, tags = {}, {};
 334         for k,v in pairs(stanza.attr) do attr[k] = v; end
 335         local new = { name = stanza.name, attr = attr, tags = tags };
 336         for i=1,#stanza do
 337                 local child = stanza[i];
 338                 if child.name then
 339                         child = _clone(child);
 340                         t_insert(tags, child);
 341                 end
 342                 t_insert(new, child);
 343         end
 344         return setmetatable(new, stanza_mt);
 345 end
 346 clone = _clone;
 347
 348 function message(attr, body)
 349         if not body then
 350                 return stanza("message", attr);
 351         else
 352                 return stanza("message", attr):tag("body"):text(body):up();
 353         end
 354 end
 355 function iq(attr)
 356         if attr and not attr.id then attr.id = new_id(); end
 357         return stanza("iq", attr or { id = new_id() });
 358 end
 359
 360 function reply(orig)
 361         return stanza(orig.name, orig.attr and { to = orig.attr.from, from = orig.attr.to, id = orig.attr.id, type = ((orig.name == "iq" and "result") or orig.attr.type) });
 362 end
 363
 364 do
 365         local xmpp_stanzas_attr = { xmlns = xmlns_stanzas };
 366         function error_reply(orig, type, condition, message)
 367                 local t = reply(orig);
 368                 t.attr.type = "error";
 369                 t:tag("error", {type = type}) --COMPAT: Some day xmlns:stanzas goes here
 370                         :tag(condition, xmpp_stanzas_attr):up();
 371                 if (message) then t:tag("text", xmpp_stanzas_attr):text(message):up(); end
 372                 return t; -- stanza ready for adding app-specific errors
 373         end
 374 end
 375
 376 function presence(attr)
 377         return stanza("presence", attr);
 378 end
 379
 380 if do_pretty_printing then
 381         local style_attrk = getstyle("yellow");
 382         local style_attrv = getstyle("red");
 383         local style_tagname = getstyle("red");
 384         local style_punc = getstyle("magenta");
 385
 386         local attr_format = " "..getstring(style_attrk, "%s")..getstring(style_punc, "=")..getstring(style_attrv, "'%s'");
 387         local top_tag_format = getstring(style_punc, "<")..getstring(style_tagname, "%s").."%s"..getstring(style_punc, ">");
 388         --local tag_format = getstring(style_punc, "<")..getstring(style_tagname, "%s").."%s"..getstring(style_punc, ">").."%s"..getstring(style_punc, "</")..getstring(style_tagname, "%s")..getstring(style_punc, ">");
 389         local tag_format = top_tag_format.."%s"..getstring(style_punc, "</")..getstring(style_tagname, "%s")..getstring(style_punc, ">");
 390         function stanza_mt.pretty_print(t)
 391                 local children_text = "";
 392                 for n, child in ipairs(t) do
 393                         if type(child) == "string" then
 394                                 children_text = children_text .. xml_escape(child);
 395                         else
 396                                 children_text = children_text .. child:pretty_print();
 397                         end
 398                 end
 399
 400                 local attr_string = "";
 401                 if t.attr then
 402                         for k, v in pairs(t.attr) do if type(k) == "string" then attr_string = attr_string .. s_format(attr_format, k, tostring(v)); end end
 403                 end
 404                 return s_format(tag_format, t.name, attr_string, children_text, t.name);
 405         end
 406
 407         function stanza_mt.pretty_top_tag(t)
 408                 local attr_string = "";
 409                 if t.attr then
 410                         for k, v in pairs(t.attr) do if type(k) == "string" then attr_string = attr_string .. s_format(attr_format, k, tostring(v)); end end
 411                 end
 412                 return s_format(top_tag_format, t.name, attr_string);
 413         end
 414 else
 415         -- Sorry, fresh out of colours for you guys ;)
 416         stanza_mt.pretty_print = stanza_mt.__tostring;
 417         stanza_mt.pretty_top_tag = stanza_mt.top_tag;
 418 end
 419
 420 return _M;