mirror of
https://github.com/maxDcb/C2Core
synced 2026-06-08 15:48:01 +00:00
Maxime dcb
108a370807
* CommandSpecs * CommandSpecs * feat(command-specs): add simple module specs * listModule * AssemblyExec * AssemblyExecTests * Minor * Inject * InjectTests * Spec modules * Folder layout * upload & download * Chisel Minidump Powershell & Script * CoffLoader DotnetExec /KerberosUseTicket PsExec PwSh ScreenShot * add artifact_filters * Minor * stabilisation * Fixes * manual test * manual test * manual test * manual test * manual test * manual test * manual test * manual test * ScreenShot png * socks5 hostname * Maj for AI * minor
109 lines
3.6 KiB
C++
109 lines
3.6 KiB
C++
#include "../Download.hpp"
|
|
#include <chrono>
|
|
#include <filesystem>
|
|
#include <functional>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <thread>
|
|
|
|
bool testDownload();
|
|
|
|
int main() {
|
|
std::cout << "[+] testDownload" << std::endl;
|
|
bool res = testDownload();
|
|
if (res) {
|
|
std::cout << "[+] Success" << std::endl;
|
|
} else {
|
|
std::cout << "[-] Failed" << std::endl;
|
|
}
|
|
return !res;
|
|
}
|
|
|
|
static bool filesEqual(const std::filesystem::path& a, const std::filesystem::path& b) {
|
|
std::ifstream f1(a, std::ios::binary);
|
|
std::ifstream f2(b, std::ios::binary);
|
|
std::string d1((std::istreambuf_iterator<char>(f1)), {});
|
|
std::string d2((std::istreambuf_iterator<char>(f2)), {});
|
|
return d1 == d2;
|
|
}
|
|
|
|
static std::filesystem::path uniqueTestTempPath(const char* prefix) {
|
|
const auto suffix = std::to_string(std::chrono::steady_clock::now().time_since_epoch().count())
|
|
+ "_" + std::to_string(std::hash<std::thread::id>{}(std::this_thread::get_id()));
|
|
return std::filesystem::temp_directory_path() / (std::string(prefix) + "_" + suffix);
|
|
}
|
|
|
|
bool testDownload() {
|
|
namespace fs = std::filesystem;
|
|
fs::path temp = uniqueTestTempPath("c2core_download_test");
|
|
fs::create_directories(temp);
|
|
bool ok = true;
|
|
|
|
// Small file
|
|
{
|
|
Download dl;
|
|
fs::path src = temp / "src.txt";
|
|
fs::path dst = temp / "dst.txt";
|
|
std::string data = "small file";
|
|
std::ofstream(src) << data;
|
|
|
|
std::vector<std::string> cmd = {"download", src.string(), dst.string()};
|
|
C2Message msg, ret;
|
|
dl.init(cmd, msg);
|
|
dl.process(msg, ret);
|
|
dl.followUp(ret);
|
|
std::cout << "small ret: " << ret.returnvalue() << " ec=" << ret.errorCode() << std::endl;
|
|
bool sub = ret.errorCode() == -1 && ret.returnvalue() == "Success" && filesEqual(src, dst);
|
|
std::cout << "small ok=" << sub << std::endl;
|
|
ok &= sub;
|
|
}
|
|
|
|
// Large file requiring multiple chunks
|
|
{
|
|
Download dl;
|
|
fs::path src = temp / "large.bin";
|
|
fs::path dst = temp / "large_copy.bin";
|
|
const size_t size = 1024 * 1024 + 500; // > CHUNK_SIZE
|
|
std::string data(size, 'A');
|
|
std::ofstream(src, std::ios::binary).write(data.data(), data.size());
|
|
|
|
std::vector<std::string> cmd = {"download", src.string(), dst.string()};
|
|
C2Message msg, ret;
|
|
dl.init(cmd, msg);
|
|
msg.set_uuid("download-uuid");
|
|
dl.process(msg, ret);
|
|
dl.followUp(ret);
|
|
std::cout << "large first ret: " << ret.returnvalue() << " ec=" << ret.errorCode() << std::endl;
|
|
while(ret.returnvalue() != "Success") {
|
|
C2Message next;
|
|
dl.recurringExec(next);
|
|
ok &= next.uuid() == "download-uuid";
|
|
ok &= next.inputfile() == src.string();
|
|
dl.followUp(next);
|
|
ret = next;
|
|
}
|
|
bool sub = filesEqual(src, dst);
|
|
std::cout << "large ok=" << sub << std::endl;
|
|
ok &= sub;
|
|
}
|
|
|
|
// Non-existent source file
|
|
{
|
|
Download dl;
|
|
fs::path src = temp / "missing.txt";
|
|
fs::path dst = temp / "missing_out.txt";
|
|
std::vector<std::string> cmd = {"download", src.string(), dst.string()};
|
|
C2Message msg, ret;
|
|
dl.init(cmd, msg);
|
|
dl.process(msg, ret);
|
|
dl.followUp(ret);
|
|
std::cout << "missing ret: " << ret.returnvalue() << " ec=" << ret.errorCode() << std::endl;
|
|
bool sub = ret.errorCode() != -1 && !fs::exists(dst);
|
|
std::cout << "missing ok=" << sub << std::endl;
|
|
ok &= sub;
|
|
}
|
|
|
|
fs::remove_all(temp);
|
|
return ok;
|
|
}
|