source: nscp/modules/LUAScript/lua_wrappers.hpp @ a8c6e93

0.4.00.4.10.4.2
Last change on this file since a8c6e93 was a8c6e93, checked in by Michael Medin <michael@…>, 16 months ago

Now "everything" (ish) works including channels, exec and query (via moderna API)
Still no protocol buffer support but not sure how to play that yet so will not be avalible in 0.4.0.

  • Property mode set to 100644
File size: 15.5 KB
Line 
1#pragma once
2
3extern "C" {
4#include <lua.h>
5#include "lauxlib.h"
6#include "lualib.h"
7}
8#include "luna.h"
9
10
11namespace lua_wrappers {
12
13
14
15        class Lua_State {
16                lua_State *L;
17        public:
18                Lua_State() : L(lua_open()) { }
19
20                ~Lua_State() {
21                        lua_close(L);
22                }
23
24                // implicitly act as a lua_State pointer
25                inline operator lua_State*() const {
26                        return L;
27                }
28                inline lua_State* get_state() const {
29                        int i = lua_gettop(L);
30                        return L;
31                }
32
33
34
35        };
36
37        class lua_wrapper {
38
39                lua_State *L;
40        public:
41
42                lua_wrapper(lua_State *L) : L(L) {}
43
44                //////////////////////////////////////////////////////////////////////////
45                /// get_xxx
46                std::wstring get_string(int pos = -1) {
47                        if (pos == -1)
48                                pos = lua_gettop(L);
49                        if (pos == 0)
50                                return _T("<EMPTY>");
51                        if (is_string(pos))
52                                return utf8::cvt<std::wstring>(lua_tostring(L, pos));
53                        if (is_number(pos))
54                                return strEx::itos(lua_tonumber(L, pos));
55                        return _T("<NOT_A_STRING>");
56                }
57                int get_int(int pos = -1) {
58                        if (pos == -1)
59                                pos = lua_gettop(L);
60                        if (pos == 0)
61                                return 0;
62                        if (is_string(pos))
63                                return strEx::stoi(utf8::cvt<std::wstring>(lua_tostring(L, pos)));
64                        if (is_number(pos))
65                                return lua_tonumber(L, pos);
66                        return 0;
67                }
68                boolean get_boolean(int pos = -1) {
69                        if (pos == -1)
70                                pos = lua_gettop(L);
71                        if (pos == 0)
72                                return false;
73                        if (is_boolean(pos))
74                                return lua_toboolean(L, pos);
75                        if (is_number(pos))
76                                return lua_tonumber(L, pos)==1;
77                        return false;
78                }
79                NSCAPI::nagiosReturn get_code(int pos = -1) {
80                        std::string str;
81                        if (pos == -1)
82                                pos = lua_gettop(L);
83                        if (pos == 0)
84                                return NSCAPI::returnUNKNOWN;
85                        switch (lua_type(L, pos)) {
86                                case LUA_TNUMBER:
87                                        return static_cast<int>(lua_tonumber(L, pos));
88                                case LUA_TSTRING:
89                                        return string_to_code(lua_tostring(L, pos));
90                                case LUA_TBOOLEAN:
91                                        return lua_toboolean(L, pos)?NSCAPI::returnOK:NSCAPI::returnCRIT;
92                        }
93                        NSC_LOG_ERROR_STD(_T("Incorrect return from script: should be error, ok, warning or unknown"));
94                        return NSCAPI::returnUNKNOWN;
95                }
96                std::list<std::wstring> get_array(const int pos = -1) {
97                        std::list<std::wstring> ret;
98                        const int len = lua_objlen(L, pos);
99                        for ( int i = 1; i <= len; ++i ) {
100                                lua_pushinteger(L, i);
101                                lua_gettable(L, -2);
102                                ret.push_back(get_string(-1));
103                                pop();
104                        }
105                        return ret;
106                }
107
108                //////////////////////////////////////////////////////////////////////////
109                /// pop_xxx
110                boolean pop_boolean() {
111                        int pos = lua_gettop(L);
112                        if (pos == 0)
113                                return false;
114                        NSCAPI::nagiosReturn ret = get_boolean(pos);
115                        lua_pop(L, 1);
116                        return ret;
117                }
118                std::wstring pop_string() {
119                        std::wstring ret;
120                        int top = lua_gettop(L);
121                        if (top == 0)
122                                return _T("<EMPTY>");
123                        ret = get_string(top);
124                        pop();
125                        return ret;
126                }
127                int pop_int() {
128                        int ret;
129                        int top = lua_gettop(L);
130                        if (top == 0)
131                                return 0;
132                        ret = get_int(top);
133                        pop();
134                        return ret;
135                }
136                NSCAPI::nagiosReturn pop_code() {
137                        int pos = lua_gettop(L);
138                        if (pos == 0)
139                                return NSCAPI::returnUNKNOWN;
140                        NSCAPI::nagiosReturn ret = get_code(pos);
141                        pop();
142                        return ret;
143                }
144                std::list<std::wstring> pop_array() {
145                        std::list<std::wstring> ret;
146                        int pos = lua_gettop(L);
147                        if (pos == 0)
148                                return ret;
149                        ret = get_array(pos);
150                        pop();
151                        return ret;
152                }
153                //////////////////////////////////////////////////////////////////////////
154                // Converters
155                NSCAPI::nagiosReturn string_to_code(std::string str) {
156                        if ((str == "critical")||(str == "crit")||(str == "error")) {
157                                return NSCAPI::returnCRIT;
158                        } else if ((str == "warning")||(str == "warn")) {
159                                return NSCAPI::returnWARN;
160                        } else if (str == "ok") {
161                                return NSCAPI::returnOK;
162                        } else if (str == "unknown") {
163                                return NSCAPI::returnUNKNOWN;
164                        }
165                        NSC_LOG_ERROR_STD(_T("Invalid code: ") + utf8::to_unicode(str));
166                        return NSCAPI::returnUNKNOWN;
167                }
168
169
170                ////////////////////////////////////////////////////////////////////////////
171                // Misc
172                inline void pop(int count = 1) {
173                        lua_pop(L, count);
174                }
175                inline int type(int pos = -1) {
176                        if (pos == -1)
177                                pos = lua_gettop(L);
178                        if (pos == 0)
179                                return LUA_TNIL;
180                        int type = lua_type(L, pos);
181                        return type;
182                }
183                std::wstring get_type_as_string(int pos = -1) {
184                        if (pos == -1)
185                                pos = lua_gettop(L);
186                        if (pos == 0)
187                                return _T("<EMPTY>");
188                        switch (lua_type(L, pos)) {
189                                case LUA_TNUMBER:
190                                        return _T("<NUMBER>");
191                                case LUA_TSTRING:
192                                        return _T("<STRING>");
193                                case LUA_TBOOLEAN:
194                                        return _T("<TABLE>");
195                                case LUA_TLIGHTUSERDATA:
196                                        return _T("<LIGHTUSERDATA>");
197                                case LUA_TTABLE:
198                                        return _T("<TABLE>");
199                        }
200                        return _T("<UNKNOWN>");
201                }
202
203                inline bool is_string(int pos = -1) {
204                        return type(pos) == LUA_TSTRING;
205                }
206                inline bool is_function(int pos = -1) {
207                        return type(pos) == LUA_TFUNCTION;
208                }
209                inline bool is_number(int pos = -1) {
210                        return type(pos) == LUA_TNUMBER;
211                }
212                inline bool is_nil(int pos = -1) {
213                        return type(pos) == LUA_TNIL;
214                }
215                inline bool is_boolean(int pos = -1) {
216                        return type(pos) == LUA_TBOOLEAN;
217                }
218                inline bool is_table(int pos = -1) {
219                        return type(pos) == LUA_TTABLE;
220                }
221
222
223                //////////////////////////////////////////////////////////////////////////
224                // push_xxx
225                void push_code(NSCAPI::nagiosReturn code) {
226                        if (code == NSCAPI::returnOK)
227                                lua_pushstring(L, strEx::wstring_to_string(_T("ok")).c_str());
228                        else if (code == NSCAPI::returnWARN)
229                                lua_pushstring(L, strEx::wstring_to_string(_T("warning")).c_str());
230                        else if (code == NSCAPI::returnCRIT)
231                                lua_pushstring(L, strEx::wstring_to_string(_T("critical")).c_str());
232                        else
233                                lua_pushstring(L, strEx::wstring_to_string(_T("unknown")).c_str());
234                }
235                void push_string(std::wstring s) {
236                        lua_pushstring(L, strEx::wstring_to_string(s).c_str());
237                }
238                void push_boolean(bool b) {
239                        lua_pushboolean(L, b?TRUE:FALSE);
240                }
241                void push_int(int b) {
242                        lua_pushinteger(L, b);
243                }
244                void push_string(std::string s) {
245                        lua_pushstring(L, s.c_str());
246                }
247                void push_array(std::list<std::wstring> &arr) {
248                        lua_createtable(L, 0, arr.size());
249                        int i=0;
250                        BOOST_FOREACH(const std::wstring &s, arr) {
251                                lua_pushnumber(L,i++);
252                                lua_pushstring(L,strEx::wstring_to_string(s).c_str());
253                                lua_settable(L,-3);
254                        }
255                }
256                inline int size() {
257                        return lua_gettop(L);
258                }
259                inline bool empty() {
260                        return size() == 0;
261                }
262                void log_stack() {
263                        int args = size();
264                        NSC_DEBUG_MSG_STD(_T("Invalid lua stack state, dumping stack"));
265                        for (int i=1;i<args+1;i++) {
266                                NSC_DEBUG_MSG_STD(get_type_as_string(i) +_T(": ") + get_string(i));
267                        }
268                }
269
270                int error(std::string s) {
271                        return luaL_error(L, s.c_str());
272                }
273
274                typedef std::pair<std::wstring,int> stack_trace;
275                stack_trace get_stack_trace(int level = 1) {
276                        lua_Debug ar;
277                        if (lua_getstack(L, level, &ar)) {  /* check function at level */
278                                lua_getinfo(L, "Sl", &ar);  /* get info about it */
279                                if (ar.currentline > 0) {  /* is there info? */
280                                        return stack_trace(utf8::cvt<std::wstring>(ar.short_src), ar.currentline);
281                                }
282                        }
283                        return stack_trace(_T("unknown"),0);
284                }
285
286                std::wstring dump_stack() {
287                        std::wstring ret;
288                        while (!empty()) {
289                                if (!ret.empty())
290                                        ret += _T(", ");
291                                ret += pop_string();
292                        }
293                        return ret;
294                }
295
296                inline void openlibs() {
297                        luaL_openlibs(L);
298                }
299
300                inline int loadfile(std::string script) {
301                        return luaL_loadfile(L, script.c_str());
302                }
303
304                int pcall(int nargs, int nresults, int errfunc) {
305                        return lua_pcall(L, nargs, nresults, errfunc);
306                }
307
308
309                std::string inline op_string(int pos, std::string def = "") {
310                        return luaL_optstring(L, pos, def.c_str());
311                }
312                std::wstring inline op_wstring(int pos, std::string def = "") {
313                        return utf8::cvt<std::wstring>(op_string(pos, def));
314                }
315                std::wstring inline op_wstring(int pos, std::wstring def) {
316                        return op_wstring(pos, utf8::cvt<std::string>(def));
317                }
318                std::string inline string(int pos) {
319                        return luaL_checkstring(L, pos);
320                }
321                std::wstring inline wstring(int pos) {
322                        return utf8::cvt<std::wstring>(string(pos));
323                }
324
325                std::list<std::wstring> inline checkarray(int pos) {
326                        luaL_checktype(L, pos, LUA_TTABLE);
327                        return get_array(pos);
328                }
329
330                boolean inline checkbool(int pos) {
331                        return lua_toboolean(L, pos);
332                }
333                int inline op_int(int pos, int def = 0) {
334                        return luaL_optinteger(L, pos, def);
335                }
336                int inline checkint(int pos) {
337                        return luaL_checkint(L, pos);
338                }
339        };
340
341        class LUAException : std::exception {
342                std::string error_;
343        public:
344                LUAException(std::wstring error) : error_(utf8::cvt<std::string>(error)) {}
345                LUAException(std::string error) : error_(error) {}
346
347                ~LUAException() throw() {}
348                const char* what() const throw() {
349                        return error_.c_str();
350                }
351                std::wstring getMessage() const {
352                        return utf8::cvt<std::wstring>(error_);
353                }
354        };
355
356        inline std::string w2s(std::wstring s) {
357                return utf8::cvt<std::string>(s);
358        }
359        inline std::wstring s2w(std::string s) {
360                return utf8::cvt<std::wstring>(s);
361        }
362
363        class lua_registry;
364        class lua_script_instance {
365                Lua_State L;
366                nscapi::core_wrapper* core;
367                int plugin_id;
368                boost::shared_ptr<lua_registry> registry;
369                std::string alias;
370                std::string script;
371
372        public:
373
374                lua_script_instance(nscapi::core_wrapper* core, int plugin_id, boost::shared_ptr<lua_registry> registry, std::string alias, std::string script)
375                        : core(core)
376                        , plugin_id(plugin_id)
377                        , registry(registry)
378                        , alias(alias)
379                        , script(script)
380                {}
381
382                int get_plugin_id() const {
383                        return plugin_id;
384                }
385                lua_State *get_lua_state() const {
386                        return L.get_state();
387                }
388                nscapi::core_wrapper* get_core() const {
389                        return core;
390                }
391                boost::shared_ptr<lua_registry> get_registry() const {
392                        return registry;
393                }
394                std::string get_script() const {
395                        return script;
396                }
397                std::string get_alias() const {
398                        return alias;
399                }
400        };
401
402        class lua_registry {
403                struct function_container {
404                        boost::shared_ptr<lua_script_instance> instance;
405                        int func_ref;
406                };
407
408                typedef std::map<std::wstring,function_container> function_map;
409                function_map functions;
410                function_map channels;
411                function_map execs;
412
413                inline lua_State * prep_function(const function_container &c) {
414                        lua_State *L = c.instance->get_lua_state();
415                        lua_rawgeti(L, LUA_REGISTRYINDEX, c.func_ref);
416                        return L;
417                }
418        public:
419
420                NSCAPI::nagiosReturn on_query(const std::wstring & target, const std::wstring & command, std::list<std::wstring> & arguments, std::wstring & message, std::wstring & perf) {
421                        function_map::iterator it = functions.find(command);
422                        if (it == functions.end())
423                                throw LUAException(_T("Invalid function: ") + command);
424                        function_container c = it->second;
425                        lua_wrapper lua(prep_function(c));
426                        lua.push_string(command);
427                        lua.push_array(arguments);
428                        if (lua.pcall(2, LUA_MULTRET, 0) != 0) {
429                                NSC_LOG_ERROR_STD(_T("Failed to handle command: ") + command + _T(": ") + lua.pop_string());
430                                return NSCAPI::returnUNKNOWN;
431                        }
432                        int arg_count = lua.size();
433                        if (arg_count > 2)
434                                perf = lua.pop_string();
435                        if (arg_count > 1)
436                                message = lua.pop_string();
437                        if (arg_count > 0)
438                                return lua.pop_code();
439                        NSC_LOG_ERROR_STD(_T("No arguments returned from script."));
440                        return NSCAPI::returnUNKNOWN;
441                }
442
443                NSCAPI::nagiosReturn on_exec(const std::wstring & command, std::list<std::wstring> & arguments, std::wstring & result) {
444                        function_map::iterator it = execs.find(command);
445                        if (it == execs.end())
446                                throw LUAException(_T("Invalid function: ") + command);
447                        function_container c = it->second;
448                        lua_wrapper lua(prep_function(c));
449                        lua.push_string(command);
450                        lua.push_array(arguments);
451                        if (lua.pcall(2, LUA_MULTRET, 0) != 0) {
452                                NSC_LOG_ERROR_STD(_T("Failed to handle command: ") + command + _T(": ") + lua.pop_string());
453                                return NSCAPI::returnUNKNOWN;
454                        }
455                        int arg_count = lua.size();
456                        if (arg_count > 1)
457                                result = lua.pop_string();
458                        if (arg_count > 0)
459                                return lua.pop_code();
460                        NSC_LOG_ERROR_STD(_T("No arguments returned from script."));
461                        return NSCAPI::returnUNKNOWN;
462                }
463                NSCAPI::nagiosReturn on_submission(const std::wstring channel, const std::wstring source, const std::wstring command, NSCAPI::nagiosReturn code, std::wstring msg, std::wstring perf) {
464                        function_map::iterator it = channels.find(channel);
465                        if (it == channels.end())
466                                throw LUAException(_T("Invalid function: ") + channel);
467                        function_container c = it->second;
468                        lua_wrapper lua(prep_function(c));
469                        lua.push_string(source);
470                        lua.push_string(command);
471                        lua.push_code(code);
472                        lua.push_string(msg);
473                        lua.push_string(perf);
474                        if (lua.pcall(5, LUA_MULTRET, 0) != 0) {
475                                NSC_LOG_ERROR_STD(_T("Failed to handle command: ") + command + _T(": ") + lua.pop_string());
476                                return NSCAPI::returnUNKNOWN;
477                        }
478                        if (lua.size() > 0)
479                                return lua.pop_code();
480                        NSC_LOG_ERROR_STD(_T("No arguments returned from script."));
481                        return NSCAPI::returnUNKNOWN;
482                }
483
484                void register_query(const std::wstring &command, boost::shared_ptr<lua_script_instance> instance, int func_ref) {
485                        function_container c;
486                        c.func_ref = func_ref;
487                        c.instance = instance;
488                        functions[command] = c;
489                }
490                void register_subscription(const std::wstring &channel, boost::shared_ptr<lua_script_instance> instance, int func_ref) {
491                        function_container c;
492                        c.func_ref = func_ref;
493                        c.instance = instance;
494                        channels[channel] = c;
495                }
496
497                void register_exec(const std::wstring &command, boost::shared_ptr<lua_script_instance> instance, int func_ref) {
498                        function_container c;
499                        c.func_ref = func_ref;
500                        c.instance = instance;
501                        execs[command] = c;
502                }
503
504                void clear() {
505                        functions.clear();
506                        execs.clear();
507                        channels.clear();
508                        // DO we need to release reference here?
509                }
510
511                bool has_command(const std::wstring & command) {
512                        return functions.find(command) != functions.end();
513                }
514                bool has_exec(const std::wstring & command) {
515                        return execs.find(command) != execs.end();
516                }
517                bool has_submit(const std::wstring &command) {
518                        return channels.find(command) != channels.end();
519                }
520        };
521
522        class lua_instance_manager {
523        public:
524                typedef boost::shared_ptr<lua_script_instance> script_instance_type;
525                typedef std::vector<script_instance_type> script_map_type;
526        private:
527                static script_map_type scripts;
528        public:
529                static void set_script(lua_State *L, script_instance_type script) {
530                        int index = 0;
531                        {
532                                // TODO: mutex lock!
533                                index = scripts.size();
534                                scripts.push_back(script);
535                        }
536                        set_registry(L, get_script_key(), index);
537                }
538                static script_instance_type get_script(lua_State *L) {
539                        int index = get_registry(L, get_script_key());
540                        {
541                                // TODO: mutex lock!
542                                if (index >= scripts.size())
543                                        throw lua_wrappers::LUAException(_T("Could not find script reference"));
544                                return scripts[index];
545                        }
546                }
547                static std::string get_script_key() {
548                        return "nscp.keys.script";
549                }
550                //////////////////////////////////////////////////////////////////////////
551                // Registry access
552                static int get_registry(lua_State *L, std::string key) {
553                        lua_pushstring(L, key.c_str());
554                        lua_gettable(L, LUA_REGISTRYINDEX);
555                        return lua_tonumber(L, -1);
556                }
557                static void set_registry(lua_State *L, std::string key, int val) {
558                        lua_pushstring(L, key.c_str());
559                        lua_pushnumber(L, val);
560                        lua_settable(L, LUA_REGISTRYINDEX);
561                }
562        };
563
564        lua_instance_manager::script_map_type lua_instance_manager::scripts;
565
566
567
568}
569
Note: See TracBrowser for help on using the repository browser.