source: nscp/modules/LUAScript/script_wrapper.hpp @ 367bf20

0.4.00.4.10.4.2stable
Last change on this file since 367bf20 was 367bf20, checked in by Michael Medin <michael@…>, 5 years ago

2008-04-03 MickeM

  • Moved COM init to "core" (from WMI module) + Added new Check command: CheckTaskSched Use like so:

CheckTaskSched +filter-exit-code==1 ShowAll MaxWarn=1 MaxCrit=1
  • Property mode set to 100644
File size: 11.3 KB
Line 
1#pragma once
2
3#include <map>
4
5extern "C" {
6#include "LUA/lua.h"
7#include "LUA/lauxlib.h"
8#include "LUA/lualib.h"
9}
10#include "luna.h"
11
12
13namespace script_wrapper {
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*() {
26                        return L;
27                }
28        };
29
30        class LUAException {
31                std::wstring error_;
32        public:
33                LUAException(std::wstring error) : error_(error) {}
34
35                std::wstring getMessage() const {
36                        return error_;
37                }
38
39        };
40
41        inline std::string w2s(std::wstring s) {
42                return strEx::wstring_to_string(s);
43        }
44        inline std::wstring s2w(std::string s) {
45                return strEx::string_to_wstring(s);
46        }
47        typedef std::pair<std::wstring,int> where_type;
48        where_type where(lua_State *L, int level = 1) {
49                lua_Debug ar;
50                if (lua_getstack(L, level, &ar)) {  /* check function at level */
51                        lua_getinfo(L, "Sl", &ar);  /* get info about it */
52                        if (ar.currentline > 0) {  /* is there info? */
53                                return where_type(s2w(ar.short_src), ar.currentline);
54                        }
55                }
56                return where_type(_T("unknown"),0);
57        }
58        std::wstring extract_string(lua_State *L) {
59                int top = lua_gettop(L);
60                if (lua_isstring(L, top))
61                        return strEx::string_to_wstring(lua_tostring( L, lua_gettop( L ) ));
62                return _T("<NOT_A_STRING>");
63        }
64        std::wstring pop_string(lua_State *L) {
65                std::wstring ret;
66                int top = lua_gettop(L);
67                if (lua_isstring(L, top))
68                        ret = strEx::string_to_wstring(lua_tostring( L, top));
69                else if (lua_isnil(L, top))
70                        ret = _T("<NIL>");
71                else if (lua_istable(L, top))
72                        ret = _T("<TABLE>");
73                else if (lua_isnumber(L, top))
74                        ret = _T("<NUMBER>");
75                else if (lua_iscfunction(L, top))
76                        ret = _T("<C-FUNCTION>");
77                else
78                        ret = _T("<UNKNOWN>");
79                lua_pop(L, 1);
80                return ret;
81        }
82        NSCAPI::nagiosReturn extract_code(lua_State *L) {
83                std::string str;
84                switch (lua_type( L, lua_gettop( L ) )) {
85        case LUA_TNUMBER:
86                return static_cast<int>(lua_tonumber(L, lua_gettop(L)));
87        case LUA_TTABLE:
88                NSC_LOG_ERROR_STD(_T("Incorect return from script: should be error, ok, warning or unknown"));
89                return NSCAPI::returnUNKNOWN;
90        case LUA_TSTRING:
91                str = lua_tostring(L, lua_gettop(L));
92                if ((str == "critical")||(str == "crit")||(str == "error")) {
93                        return NSCAPI::returnCRIT;
94                } else if ((str == "warning")||(str == "warn")) {
95                        return NSCAPI::returnWARN;
96                } else if (str == "ok") {
97                        return NSCAPI::returnOK;
98                } else if (str == "unknown") {
99                        return NSCAPI::returnUNKNOWN;
100                } else {
101                        NSC_LOG_ERROR_STD(_T("Incorect return from script: should be ok, warning, critical or unknown not: ") + strEx::string_to_wstring(str) );
102                        return NSCAPI::returnUNKNOWN;
103                }
104        case LUA_TBOOLEAN:
105                return lua_toboolean( L, lua_gettop( L ) )?NSCAPI::returnOK:NSCAPI::returnCRIT;
106                }
107                NSC_LOG_ERROR_STD(_T("Incorect return from script: should be error, ok, warning or unknown"));
108                return NSCAPI::returnUNKNOWN;
109        }
110        void push_code(lua_State *L, NSCAPI::nagiosReturn  code) {
111                if (code == NSCAPI::returnOK)
112                        lua_pushstring(L, strEx::wstring_to_string(_T("ok")).c_str());
113                else if (code == NSCAPI::returnWARN)
114                        lua_pushstring(L, strEx::wstring_to_string(_T("warning")).c_str());
115                else if (code == NSCAPI::returnCRIT)
116                        lua_pushstring(L, strEx::wstring_to_string(_T("critical")).c_str());
117                else
118                lua_pushstring(L, strEx::wstring_to_string(_T("unknown")).c_str());
119        }
120        void push_string(lua_State *L, std::wstring s) {
121                lua_pushstring(L, strEx::wstring_to_string(s).c_str());
122        }
123
124        class lua_script;
125        class lua_handler {
126        public:
127                virtual void register_command(lua_script* script, std::wstring command, std::wstring function) = 0;
128
129        };
130        class lua_manager {
131                typedef std::map<double,lua_handler*> handler_type;
132                typedef std::map<double,lua_script*> script_type;
133                static handler_type handlers;
134                static script_type scripts;
135                static double last_value;
136                static char handler_key[];
137                static char script_key[];
138        public:
139                static lua_handler* get_handler(lua_State *L) {
140                        handler_type::const_iterator cit = handlers.find(get_id(L, handler_key));
141                        if (cit == handlers.end())
142                                throw LUAException(_T("Could not find handler reference"));
143                        return (*cit).second;
144                }
145                static void set_handler(lua_State *L, lua_handler* handler) {
146                        double id = get_id(L, handler_key);
147                        handlers[id] = handler;
148                }
149                static lua_script* get_script(lua_State *L) {
150                        script_type::const_iterator cit = scripts.find(get_id(L, script_key));
151                        if (cit == scripts.end())
152                                throw LUAException(_T("Could not find script reference"));
153                        return (*cit).second;
154                }
155                static void set_script(lua_State *L, lua_script* script) {
156                        double id = get_id(L, script_key);
157                        scripts[id] = script;
158                }
159                static double get_id(lua_State *L, char *key) {
160                        /* retrieve a number */
161                        lua_pushstring(L, key);
162                        //lua_pushlightuserdata(L, (void*)&key);  /* push address */
163                        lua_gettable(L, LUA_REGISTRYINDEX);  /* retrieve value */
164                        double v = 0;
165                        v = lua_tonumber(L, -1);  /* convert to number */
166                        lua_pop(L,1);
167                        if (v <= 0) {
168                                v = ++last_value;
169                                lua_pushstring(L, key);
170                                //lua_pushlightuserdata(L, reinterpret_cast<void*>(&key));  /* push address */
171                                lua_pushnumber(L, v);  /* push value */
172                                /* registry[&Key] = myNumber */
173                                lua_settable(L, LUA_REGISTRYINDEX);
174                        }
175                        return v;
176
177                }
178
179        };
180        class nsclient_wrapper {
181        public:
182
183                static int execute (lua_State *L) {
184                        try {
185                                int nargs = lua_gettop( L );
186                                if (nargs == 0) {
187                                        return luaL_error(L, "nscp.execute requires atleast 1 argument!");
188                                }
189                                unsigned int argLen = nargs-1;
190                                arrayBuffer::arrayBuffer arguments = arrayBuffer::createArrayBuffer(argLen);
191                                for (unsigned int i=argLen;i>0;i--) {
192                                        std::wstring arg = extract_string(L);
193                                        arrayBuffer::set(arguments, argLen, i-1, arg);
194                                        lua_pop(L, 1);
195                                }
196                                std::wstring command = extract_string(L);
197                                lua_pop(L, 1);
198                                std::wstring msg;
199                                std::wstring perf;
200                                NSCAPI::nagiosReturn ret = NSCModuleHelper::InjectCommand(command.c_str(), argLen, arguments, msg, perf);
201                                push_code(L, ret);
202                                lua_pushstring(L, strEx::wstring_to_string(msg).c_str());
203                                lua_pushstring(L, strEx::wstring_to_string(perf).c_str());
204                                return 3;
205                        } catch (...) {
206                                return luaL_error(L, "Unknown exception in: nscp.execute");
207                        }
208                }
209
210                static int register_command(lua_State *L) {
211                        try {
212                                lua_handler *handler = lua_manager::get_handler(L);
213                                lua_script *script = lua_manager::get_script(L);
214                                int nargs = lua_gettop( L );
215                                if (nargs != 2)
216                                        return luaL_error(L, "Incorrect syntax: nscp.register(<key>, <function>);");
217                                handler->register_command(script, pop_string(L), pop_string(L));
218                                return 0;
219                        } catch (LUAException e) {
220                                return luaL_error(L, std::string("Error in nscp.register: " + w2s(e.getMessage())).c_str());
221                        } catch (...) {
222                                return luaL_error(L, "Unknown exception in: nscp.register");
223                        }
224                }
225
226                static int getSetting (lua_State *L) {
227                        int nargs = lua_gettop( L );
228                        if (nargs < 2 || nargs > 3)
229                                return luaL_error(L, "Incorrect syntax: nscp.getSetting(<section>, <key>[, <default value>]);");
230                        std::wstring v;
231                        if (nargs > 2)
232                                v = pop_string(L);
233                        std::wstring k = pop_string(L);
234                        std::wstring s = pop_string(L);
235                        push_string(L, NSCModuleHelper::getSettingsString(s, k, v));
236                        return 1;
237                }
238                static int getSection (lua_State *L) {
239                        NSC_DEBUG_MSG_STD(_T("LUA::setSettings"));
240                        return 0;
241                }
242                static int info (lua_State *L) {
243                        return log_any(L, NSCAPI::log);
244                }
245                static int error (lua_State *L) {
246                        return log_any(L, NSCAPI::error);
247                }
248                static int log_any(lua_State *L, int mode) {
249                        where_type w = where(L);
250                        int nargs = lua_gettop( L );
251                        std::wstring str;
252                        for (int i=0;i<nargs;i++) {
253                                str += pop_string(L);
254                        }
255                        NSCModuleHelper::Message(mode, w.first, w.second, str);
256                        return 0;
257                }
258
259                static const luaL_Reg my_funcs[];
260
261                static int luaopen(lua_State *L) {
262                        luaL_register(L, "nscp", my_funcs);
263                        return 1;
264                }
265
266
267        };
268        const luaL_Reg nsclient_wrapper::my_funcs[] = {
269                {"execute", execute},
270                {"info", info},
271                {"print", info},
272                {"error", error},
273                {"register", register_command},
274                {"getSetting", getSetting},
275                {"getSection", getSection},
276                {NULL, NULL}
277        };
278
279        lua_manager::handler_type lua_manager::handlers;
280        lua_manager::script_type lua_manager::scripts;
281        double lua_manager::last_value = 0;
282        char lua_manager::handler_key[] = "registry.key.handler";
283        char lua_manager::script_key[] = "registry.key.script";
284
285        class lua_script {
286                Lua_State L;
287                std::wstring script_;
288        public:
289                lua_script(const std::wstring file) : script_(file) {
290                        load();
291                }
292                void load() {
293                        luaL_openlibs(L);
294                        nsclient_wrapper::luaopen(L);
295                        //Luna<Account>::Register(L);
296                        //lua_register(L, "register_command", register_command);
297
298                        if (luaL_loadfile(L, strEx::wstring_to_string(script_).c_str()) != 0) {
299                                throw LUAException(_T("Failed to load script: ") + script_ + _T(": ") + s2w(lua_tostring(L, -1)));
300                        }
301
302                }
303                std::wstring get_script() const {
304                        return script_;
305                }
306                void unload() {}
307                void reload(lua_handler *handler) {
308                        unload();
309                        load();
310                        pre_load(handler);
311                }
312                void pre_load(lua_handler *handler) {
313                        lua_manager::set_handler(L, handler);
314                        lua_manager::set_script(L, this);
315                        if (lua_pcall(L, 0, 0, 0) != 0) {
316                                throw LUAException(_T("Failed to parse script: ") + script_ + _T(": ") + s2w(lua_tostring(L, -1)));
317                        }
318                }
319
320
321                NSCAPI::nagiosReturn extract_return(Lua_State &L, int arg_count,  std::wstring &message, std::wstring &perf) {
322                        // code, message, performance data
323                        if (arg_count > 3) {
324                                NSC_LOG_ERROR_STD(_T("Too many arguments return from script (only using last 3)"));
325                                lua_pop(L, arg_count-3);
326                        }
327                        if (arg_count > 2) {
328                                perf = extract_string(L);
329                                lua_pop( L, 1 );
330                        }
331                        if (arg_count > 1) {
332                                message = extract_string(L);
333                                lua_pop( L, 1 );
334                        }
335                        if (arg_count > 0) {
336                                int ret = extract_code(L);
337                                lua_pop( L, 1 );
338                                return ret;
339                        }
340                        NSC_LOG_ERROR_STD(_T("No arguments returned from script."));
341                        return NSCAPI::returnUNKNOWN;
342                }
343
344                NSCAPI::nagiosReturn handleCommand(lua_handler *handler, std::wstring function, strEx::blindstr command, const unsigned int argLen, TCHAR **char_args, std::wstring &msg, std::wstring &perf) {
345                        lua_manager::set_handler(L, handler);
346                        lua_manager::set_script(L, this);
347                        int nargs = lua_gettop( L );
348                        lua_getglobal(L, w2s(function).c_str());
349                        if (!lua_isfunction(L, -1)) {
350                                lua_pop(L, 1); // remove function from LUA stack
351                                throw LUAException(_T("Failed to run script: ") + script_ + _T(": Function not found: handle"));
352                        }
353                        std::wstring cmd = command.c_str();
354                        lua_pushstring(L, w2s(cmd).c_str());
355
356                        lua_createtable(L, 0, argLen);
357                        for (int i=0;i<argLen;i++) {
358                                lua_pushnumber(L,i+1);
359                                lua_pushstring(L,strEx::wstring_to_string(char_args[i]).c_str());
360                                lua_settable(L,-3);
361                        }
362
363                        if (lua_pcall(L, 2, LUA_MULTRET, 0) != 0) {
364                                std::wstring err = strEx::string_to_wstring(lua_tostring(L, -1));
365                                NSC_LOG_ERROR_STD(_T("Failed to call main function in script: ") + script_ + _T(": ") + err);
366                                lua_pop(L, 1); // remove error message
367                                return NSCAPI::returnUNKNOWN;
368                        }
369                        return extract_return(L, lua_gettop( L )-nargs, msg, perf);
370                }
371        };
372
373
374
375
376
377}
Note: See TracBrowser for help on using the repository browser.