Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Getting rid of huge switch statement #61

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

paul-axe
Copy link

@paul-axe paul-axe commented Oct 7, 2024

Hi. I am planning to add dynamic loading features to poseidon, however it seems to be impossible in current implementation because of this huge switch statement. This PR is trying to change this switch statement to a map of functions.

Feel free to comment with your suggestions on improvements if I missed a point somewhere

@its-a-feature
Copy link
Contributor

Nice! That mirrors the style for the adding in the functions for registration on the Mythic piece more with the init() style declarations. I'll have to pull it down and do some testing to make sure there's no issues with that mod for the prompt command, but looks cool!

What is it you're looking to do for dynamic loading?

@paul-axe
Copy link
Author

paul-axe commented Oct 7, 2024

I often faced the situation when I need some particular code (e.g. chmod) to execute on some particular agent, but want to avoid using pty. And the process of running new instance of poseidon with added specific functions can be tricky, noisy or even can crash the service because of the nature of the initial vulnerability. In that case it would be great to be able to load code into poseidon dynamically.

Currently i am working on API to dynamically download shared object from external service to in-memory file and load the code from it. Currently the main issue is that structs.Task is not exposed from the poseidon module, so I cannot easily build plugin which will add new features to the agent

@its-a-feature
Copy link
Contributor

Interesting. I didn't think Go played well with dynamically loading new functionality, so I'm super interested to see where this goes!

@paul-axe
Copy link
Author

paul-axe commented Oct 7, 2024

By the way, I found why i cannot import poseidon structs from external modules.

I trying to import structs definitions from https://github.com/MythicAgents/poseidon/blob/master/Payload_Type/poseidon/poseidon/agent_code/pkg/utils/structs/definitions.go in external golang app, but cannot do it due to some modules issue.

  • github.com/MythicAgents/poseidon/Payload_Type/poseidon - this module contains invalid module name it go.mod
  • github.com/MythicAgents/poseidon/Payload_Type/poseidon/poseidon/agent_code - this module also contains invalid module name (it declares github.com/MythicAgents/poseidon/Payload_Type/poseidon/agent_code) instead.

Should I open another PR in order to fix at least second go.mod?

@its-a-feature
Copy link
Contributor

how is it you're trying to import it? Neither were made with the intent to be libraries for 3rd party programs to import, so maybe that's why you're running into issues.

There's two separate Go projects - the one at Payload_Type/poseidon is the component of the project that connects up to Mythic and syncs data so you can see it in the Mythic UI. The one at Payload_Type/poseidon/poseidon/agent_code is the actual agent code itself. They're kept separate because the first one gets compiled when the container is created and run when the container starts, but the second one is dynamically compiled when you task a build in the UI.

@paul-axe
Copy link
Author

paul-axe commented Oct 7, 2024

Will changing the module name for Payload_Type/poseidon/poseidon/agent_code break anything?

As I said, I want to implement dynamic loaded modules, I think it will be great if these modules will not be a part of poseidon repo, but it requires to re-use some part of poseidon code, for example structs.Task in order to follow the same architecture of task handlers as we have in poseidon already. Of course it can be done with set of wrappers on the loader side, but it doesnt seems like a best solution to me.

Here is the example of the dynamically loaded task handler I want to make

package main

import (
        "github.com/MythicAgents/poseidon/Payload_Type/poseidon/agent_code"
)

func Run(structs.Task) {
       // handler logic here
}

func OnLoad(registerFunc func(string, func(structs.Task))) bool {
        registerFunc("funcName", Run)
}

So basically we are adding only OnLoad handler (which is a part of boilerplate), other code looks like ordinary task handler. The second difference is that we need to compile it in a different way. And here is where the issue comes. If we want to reuse definition of structs.Task we need to import it using go get, but the issue is that module name is not correct for import from 3rd party apps.

go get -u github.com/MythicAgents/poseidon/Payload_Type/poseidon/poseidon/agent_code
go: github.com/MythicAgents/poseidon/Payload_Type/poseidon/poseidon/agent_code@upgrade (v0.0.0-20241007155010-93d975d13ad5) requires github.com/MythicAgents/poseidon/Payload_Type/poseidon/poseidon/[email protected]: parsing go.mod:
        module declares its path as: github.com/MythicAgents/poseidon/Payload_Type/poseidon/agent_code
                but was required as: github.com/MythicAgents/poseidon/Payload_Type/poseidon/poseidon/agent_code

However, I am not sure if that is the only thing which will not allow us to reuse the part of poseidon code in dynamic modules, not really good at golang to be honest :)

@its-a-feature
Copy link
Contributor

Ah I see what you mean about the pathing now. I don't think it'll be an issue, but by all means change the path and see if it breaks anything :)

There is one other issue though. You want to create funcName externally, dynamically compile it, and load it into Poseidon. However, Mythic won't know of funcName being a valid task to send down to Poseidon, so you won't be able to issue it if you do things this way.

@paul-axe
Copy link
Author

paul-axe commented Oct 7, 2024

Regarding the Mythic side, I am planning to follow the same aproach as Apollo does here https://github.com/MythicAgents/Apollo/blob/d3e58d65be1350789ff2268998cb96356ac454bf/Payload_Type/apollo/apollo/mythic/agent_functions/load.py#L204

Not sure if I will be able to do that, but at least i am ready to try :)

@its-a-feature
Copy link
Contributor

That'll only work for commands that are already registered with Mythic. Apollo allows you to build an agent with only a subset of commands, then at run-time dynamically load commands that you chose not to include initially. It doesn't support adding completely new commands like that at run-time.

@paul-axe
Copy link
Author

paul-axe commented Oct 7, 2024

Hmm, thank you for pointing this out. Will do some more research then

@its-a-feature
Copy link
Contributor

There's still options, but that way won't work. We can also potentially add in new things to support something like this, but just gotta brainstorm how it might work. One stop-gap that you might be able to do in the meantime:

  1. create a new package in poseidon, dynamic_library
  2. create a new command, dynamic_load, in dynamic_library that takes in your custom compiled command and registers it internally
  3. create a new command, dynamic_exec, in dynamic_library that takes in the name of the function you want to execute and looks it up internally, and executes it.
    Your commands would essentially only have the ability for generic arguments (i.e. always just a single string), but that would allow you to at least move on with your testing and try to get the dynamic compilation, dynamic loading, and dynamic execution pieces worked out. We can then work on a more elegant solution

@paul-axe paul-axe force-pushed the master branch 2 times, most recently from d5b7965 to 2f2cf48 Compare October 12, 2024 09:28
@paul-axe
Copy link
Author

Well, after some research it turned out that golang officially doesnt provide any compatibility between different toolchain versions and even between builds using the same toolchain but the different source code. It means that there is always a chance that main program will crash if we will try to dynamically load external module just because compiler decided to organize structs in a different way. Considering this, my idea of dynamic module loading doesnt feel reliable enough, however I still think that getting rid of the huge switch statement is a good improvement :)

@paul-axe
Copy link
Author

However I realized that we can use C pointers to communicate between main process and loaded libraries. It will require a bit of boilerplate, manual memory management and structs serialization, but should work well.

So the current idea is to call module functions as following:

import (
    "unsafe"
    "os"
)

/*
#cgo LDFLAGS: -ldl

#include <stdlib.h>
#include <dlfcn.h>

typedef char* (*wrapper_type)(char*); // function pointer type

char* wrapper(void* f, char* s) { // wrapper function
    return ((wrapper_type) f)(s);
}
*/
import "C"

func CallLibraryFunc(library, fname, fargs string) {
    // fname is function name, fargs is serialized (JSON?) arguments
    
    libname := C.CString(library)
    defer C.free(unsafe.Pointer(libname))
    handle := C.dlopen(libname, C.RTLD_LAZY)
    if handle == nil {
        panic("Cannot open lib")
    }

    sym := C.CString(fname)
    defer C.free(unsafe.Pointer(sym))
    p := C.dlsym(handle, sym)
    if p == nil {
        panic("Cannot find symbol")
    }

    args :=  C.CString(fargs)
    defer C.free(unsafe.Pointer(args))
    cres := C.wrapper(p, args)

    result := C.GoString(cres)
    defer C.free(unsafe.Pointer(cres))
    // serialized result processing here
}

And the module part

package main

import (
    "C"
    "fmt"
)

//export run
func run(args *C.char) *C.char  {
    fmt.Printf("In library with %s\n", C.GoString(args))
    return C.CString("result")
}

func main(){}

Basically, on the module side we need only to boilerplate communication protocol, which can be provided as a reusable package I believe

@its-a-feature
Copy link
Contributor

That's pretty similar to what we do already for the execute_library command:
https://github.com/MythicAgents/poseidon/blob/master/Payload_Type/poseidon/poseidon/agent_code/execute_library/execute_library_darwin.m#L6
It goes one step further though and allows you to have argc and argv like a normal executable's main function.

The downside to this command and the one you're describing is that they're non streaming. So, say you wanted to run a long-running task with this, you'd only get output when the entire thing finishes.

@paul-axe
Copy link
Author

paul-axe commented Oct 18, 2024

We probably can fix this downside using cgo.Handle to pass complex variables to the C context and then back to the Go context. However, I realized that if we compiling dynamic library with -buildmode=c-shared it brings a full go runtime with it. Besides the high memorey consumption (every module will run its own go runtime) it will be also impossible to use anything from other go runtime directly.

However, if we dont bother about memory consumption, it seems to be possible to implement some kind of C middleware which will provide channel between different Go runtimes.

Writing dynamic modules in C should also work fine, btw :)

Here is the example how this channel can look like.
Main process:

package main
import (
    "fmt"
    "unsafe"
    "os"
)

/*
#cgo LDFLAGS: -ldl -Wl,--allow-multiple-definition

#include <stdlib.h>
#include <stdio.h>
#include <dlfcn.h>

typedef char* (*wrapper_type)(void *);

static inline char* wrapper(void* f, void *h) {
    return ((wrapper_type) f)(h);
}

void my_go_func(int);

void do_print(int a) {
    my_go_func(a);
}
*/
import "C"


var c chan int


//export my_go_func
func my_go_func(arg C.int) {
    fmt.Println("my_go_func called")
    c <- int(arg)
    fmt.Println("int sent to chan")
}

func main() {
    c = make(chan int, 10)
    name := os.Args[1]
    fmt.Printf("Trying to use %s\n", name)

    libname := C.CString(name)
    defer C.free(unsafe.Pointer(libname))
    handle := C.dlopen(libname, C.RTLD_LAZY)
    if handle == nil {
        panic("Cannot open lib")
    }

    init_sym := C.CString("module_init")
    defer C.free(unsafe.Pointer(init_sym))
    init := C.dlsym(handle, init_sym)
    if init == nil {
        panic("Cannot find symbol")
    }

    C.wrapper(init, C.do_print)
}

and the module part

package main

import (
    "C"
    "fmt"
    "unsafe"
)

/*
typedef char* (*wrapper_type)(int );
static inline char* wrapper(void* f, int h) {
    return ((wrapper_type) f)(h);
}

*/
import "C"

//export module_init
func module_init(cb unsafe.Pointer) {
    C.wrapper(cb, C.int(1337))
}

func main(){}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants