diff --git a/Makefile b/Makefile index 8dca1c1..2cbec50 100644 --- a/Makefile +++ b/Makefile @@ -30,13 +30,13 @@ flight_profiler_agent.${SHARED_LIB_SUFFIX}: @mv build/libfrida-gum.* build/lib/ @echo "compiling flight_profiler_agent.${SHARED_LIB_SUFFIX}" @${CC} ${CFLAGS} ${LDFLAGS} -I${PY_HEADER_PATH} -Ibuild/include -Icsrc \ - csrc/code_inject.cpp csrc/frida_profiler.cpp \ + csrc/profiler_attach.cpp csrc/frida_profiler.cpp \ csrc/time_util.cpp csrc/symbol_util.cpp csrc/python_util.cpp \ csrc/py_gil_intercept.cpp csrc/py_gil_stat.cpp csrc/stack/py_stack.cpp \ -o build/lib/flight_profiler_agent.${SHARED_LIB_SUFFIX} -Lbuild/lib -lfrida-gum -ldl @if [ "$(IS_DARWIN)" != "Darwin" ]; then \ - $(CC) $(INJECT_CFLAGS) -Icsrc/inject/ -o build/lib/inject csrc/inject/ProcessTracer.cpp csrc/inject/ProcessUtils.cpp csrc/inject/LibraryInjector.cpp csrc/inject/inject.cpp -ldl;\ - cp build/lib/inject flight_profiler/lib/inject;\ + $(CC) $(INJECT_CFLAGS) -Icsrc/attach/ -o build/lib/attach csrc/attach/ProcessTracer.cpp csrc/attach/ProcessUtils.cpp csrc/attach/AttachAgent.cpp csrc/attach/attach.cpp -ldl;\ + cp build/lib/attach flight_profiler/lib/attach;\ fi @cp build/lib/flight_profiler_agent.${SHARED_LIB_SUFFIX} flight_profiler/lib/flight_profiler_agent.${SHARED_LIB_SUFFIX} diff --git a/csrc/inject/LibraryInjector.cpp b/csrc/attach/AttachAgent.cpp similarity index 62% rename from csrc/inject/LibraryInjector.cpp rename to csrc/attach/AttachAgent.cpp index 9871fe3..7ca2d43 100644 --- a/csrc/inject/LibraryInjector.cpp +++ b/csrc/attach/AttachAgent.cpp @@ -1,4 +1,4 @@ -#include "LibraryInjector.h" +#include "AttachAgent.h" #include #include #include @@ -7,11 +7,11 @@ #include #include -// Assembly code for injecting shared library +// Assembly code for loading shared library into target process // This is kept as a C function because it needs to be position-independent extern "C" { /** - * @brief Injects a shared library into a target process using assembly code + * @brief Loads a shared library into a target process using assembly code * * This function performs the following steps: * 1. Saves addresses of free() and __libc_dlopen_mode() on the stack @@ -23,19 +23,19 @@ extern "C" { * @param free_function_address Address of free function in target process * @param dlopen_function_address Address of dlopen function in target process */ -void inject_shared_library(long malloc_function_address, - long free_function_address, - long dlopen_function_address, - long library_path_length); +void load_shared_library(long malloc_function_address, + long free_function_address, + long dlopen_function_address, + long library_path_length); /** - * @brief Marks the end of inject_shared_library function for size calculation + * @brief Marks the end of load_shared_library function for size calculation */ -void inject_shared_library_end(); +void load_shared_library_end(); } /** - * @brief Injects a shared library into a target process using optimized + * @brief Loads a shared library into a target process using optimized * assembly code * * This function performs the following steps: @@ -50,8 +50,8 @@ void inject_shared_library_end(); * @param free_addr Address of free function in target process * @param dlopen_addr Address of dlopen function in target process */ -void inject_shared_library(long malloc_addr, long free_addr, long dlopen_addr, - long library_path_length) { +void load_shared_library(long malloc_addr, long free_addr, long dlopen_addr, + long library_path_length) { // Optimized assembly code with macro usage for simplified operations asm( // Efficient stack alignment and register preservation @@ -85,58 +85,58 @@ void inject_shared_library(long malloc_addr, long free_addr, long dlopen_addr, } /** - * @brief Marks the end of inject_shared_library function for size calculation + * @brief Marks the end of load_shared_library function for size calculation * - * This function's only purpose is to be contiguous to inject_shared_library(), + * This function's only purpose is to be contiguous to load_shared_library(), * so that we can use its address to more precisely figure out how long - * inject_shared_library() is. + * load_shared_library() is. */ -void inject_shared_library_end() { +void load_shared_library_end() { // Intentionally empty - used only for calculating function size } /** - * @brief Constructs a LibraryInjector instance + * @brief Constructs an AttachAgent instance * - * Initializes the LibraryInjector with the target process identifier and + * Initializes the AttachAgent with the target process identifier and * library file path. Also initializes the ProcessTracer for the target process. * - * @param target_process_id PID of the target process to inject - * @param shared_library_file_path File path of the shared library to inject + * @param target_process_id PID of the target process to attach + * @param shared_library_file_path File path of the shared library to load */ -LibraryInjector::LibraryInjector(pid_t target_process_id, - const std::string &shared_library_file_path, - bool debug_mode) +AttachAgent::AttachAgent(pid_t target_process_id, + const std::string &shared_library_file_path, + bool debug_mode) : target_process_id_(target_process_id), library_file_path_(shared_library_file_path), process_tracer_(target_process_id, debug_mode) {} /** - * @brief Cleans up resources used by the LibraryInjector + * @brief Cleans up resources used by the AttachAgent * * The ProcessTracer destructor will handle cleanup of any attached processes. */ -LibraryInjector::~LibraryInjector() {} +AttachAgent::~AttachAgent() {} /** - * @brief Performs the complete library injection process + * @brief Performs the complete agent attach process * - * This function performs the complete injection process: - * 1. Initializes the injection environment by attaching to the process and + * This function performs the complete attach process: + * 1. Initializes the attach environment by attaching to the process and * getting registers * 2. Resolves necessary function addresses in the target process - * 3. Sets up registers for the injection - * 4. Orchestrates the injection sequence + * 3. Sets up registers for the attach + * 4. Orchestrates the attach sequence * - * @return true if injection was successful, false otherwise + * @return ExitCode indicating success or failure */ -ExitCode LibraryInjector::performInjection() { - long code_injection_address = 0; +ExitCode AttachAgent::performAttach() { + long code_attach_address = 0; REG_TYPE original_registers, working_registers; - // Initialize injection environment - ExitCode initialize_code = initializeInjectionEnvironment( - code_injection_address, &original_registers, &working_registers); + // Initialize attach environment + ExitCode initialize_code = initializeAttachEnvironment( + code_attach_address, &original_registers, &working_registers); if (initialize_code != ExitCode::SUCCESS) { if (initialize_code == ExitCode::GET_REGISTERS_AFTER_ATTACH_FAILED) { process_tracer_.detach(); @@ -193,35 +193,35 @@ ExitCode LibraryInjector::performInjection() { return ExitCode::SET_INJECTED_SHELLCODE_REGISTERS_FAILED; } - // Orchestrate injection sequence - ExitCode injection_result = orchestrateInjectionSequence( - code_injection_address, target_malloc_function_address, + // Orchestrate attach sequence + ExitCode attach_result = orchestrateAttachSequence( + code_attach_address, target_malloc_function_address, target_free_function_address, target_dlopen_function_address, library_file_path_.length() + 1, &original_registers); - return injection_result; + return attach_result; } /** - * @brief Initialize injection environment by attaching to the process and + * @brief Initialize attach environment by attaching to the process and * setting up registers * * This function: * 1. Attaches to the target process * 2. Gets the current register state - * 3. Finds a suitable memory address for code injection - * 4. Sets up registers for the injection + * 3. Finds a suitable memory address for code loading + * 4. Sets up registers for the attach * - * @param code_injection_address Reference to store the address where code will - * be injected + * @param code_attach_address Reference to store the address where code will + * be loaded * @param original_registers Pointer to store the original register state * @param working_registers Pointer to store the modified register state - * @return true if initialization was successful, false otherwise + * @return ExitCode indicating success or failure */ ExitCode -LibraryInjector::initializeInjectionEnvironment(long &code_injection_address, - REG_TYPE *original_registers, - REG_TYPE *working_registers) { +AttachAgent::initializeAttachEnvironment(long &code_attach_address, + REG_TYPE *original_registers, + REG_TYPE *working_registers) { // Attach to process if (!process_tracer_.attach()) { return ExitCode::ATTACH_FAILED; @@ -239,29 +239,29 @@ LibraryInjector::initializeInjectionEnvironment(long &code_injection_address, // The findFreeMemoryAddress function returns the END of the first executable // region minus a safe offset, placing shellcode in the alignment padding area // that is typically unused but still has execute permissions. - code_injection_address = + code_attach_address = ProcessUtils::findFreeMemoryAddress(target_process_id_) + 8; if (process_tracer_.isDebugMode()) { - std::cout << "[DEBUG] PyFlightProfiler: Using injection address at 0x" - << std::hex << code_injection_address << std::dec << std::endl; + std::cout << "[DEBUG] PyFlightProfiler: Using attach address at 0x" + << std::hex << code_attach_address << std::dec << std::endl; } - // Set the target's rip to the injection address + // Set the target's rip to the attach address // Advance by 2 bytes because rip gets incremented by the size of the current // instruction - working_registers->rip = code_injection_address + 2; + working_registers->rip = code_attach_address + 2; return ExitCode::SUCCESS; } /** - * @brief Create shellcode payload for library injection + * @brief Create shellcode payload for library loading * * This function: - * 1. Calculates the size of the inject_shared_library function + * 1. Calculates the size of the load_shared_library function * 2. Locates the return instruction offset * 3. Creates a buffer with NOP padding - * 4. Copies the injection function code to the buffer + * 4. Copies the library loading function code to the buffer * 5. Overwrites the return instruction with an INT 3 breakpoint * * @param payload_size Reference to store the size of the generated shellcode @@ -270,31 +270,31 @@ LibraryInjector::initializeInjectionEnvironment(long &code_injection_address, * @return Vector containing the generated shellcode */ std::vector -LibraryInjector::createShellcodePayload(size_t &payload_size, - intptr_t &return_instruction_offset) { - // Figure out the size of inject_shared_library() so we know how big of a +AttachAgent::createShellcodePayload(size_t &payload_size, + intptr_t &return_instruction_offset) { + // Figure out the size of load_shared_library() so we know how big of a // buffer to allocate. payload_size = - (intptr_t)inject_shared_library_end - (intptr_t)inject_shared_library + 2; + (intptr_t)load_shared_library_end - (intptr_t)load_shared_library + 2; // Also figure out where the RET instruction at the end of - // inject_shared_library() lies so that we can overwrite it with an INT 3 in + // load_shared_library() lies so that we can overwrite it with an INT 3 in // order to break back into the target process. return_instruction_offset = (intptr_t)ProcessUtils::locateReturnInstruction( - (void *)inject_shared_library_end) - - (intptr_t)inject_shared_library; + (void *)load_shared_library_end) - + (intptr_t)load_shared_library; - // Set up a buffer to hold the code we're going to inject into the target + // Set up a buffer to hold the code we're going to load into the target // process. std::vector shellcode_payload(payload_size, 0); shellcode_payload[0] = - 0x90; // fill with NOP, because when inject a process stuck in syscall, + 0x90; // fill with NOP, because when attach a process stuck in syscall, shellcode_payload[1] = 0x90; // rip will normally decrease by 2, so fill top // two bytes with nop instruction - // Copy the code of inject_shared_library() to a buffer. - memcpy(shellcode_payload.data() + 2, (void *)inject_shared_library, + // Copy the code of load_shared_library() to a buffer. + memcpy(shellcode_payload.data() + 2, (void *)load_shared_library, payload_size - 3); // Overwrite the RET instruction with an INT 3. @@ -304,28 +304,28 @@ LibraryInjector::createShellcodePayload(size_t &payload_size, } /** - * @brief Orchestrate the injection sequence + * @brief Orchestrate the attach sequence * - * This function performs the complete injection sequence: + * This function performs the complete attach sequence: * 1. Creates shellcode payload - * 2. Backs up original data at the injection address + * 2. Backs up original data at the attach address * 3. Deploys the shellcode - * 4. Executes the injected code + * 4. Executes the loaded code * 5. Handles the malloc() call and library path copying * 6. Calls __libc_dlopen_mode to load the library * 7. Frees the allocated buffer - * 8. Confirms the injection success + * 8. Confirms the attach success * - * @param injection_address Address where the shellcode will be injected + * @param attach_address Address where the shellcode will be loaded * @param malloc_function_address Address of malloc function in target process * @param free_function_address Address of free function in target process * @param dlopen_function_address Address of dlopen function in target process * @param library_path_string_length Length of the library path string * @param initial_registers Pointer to the original register state - * @return true if injection was successful, false otherwise + * @return ExitCode indicating success or failure */ -ExitCode LibraryInjector::orchestrateInjectionSequence( - long injection_address, long malloc_function_address, +ExitCode AttachAgent::orchestrateAttachSequence( + long attach_address, long malloc_function_address, long free_function_address, long dlopen_function_address, int library_path_string_length, REG_TYPE *initial_registers) { // Create shellcode payload @@ -334,49 +334,44 @@ ExitCode LibraryInjector::orchestrateInjectionSequence( std::vector shellcode_payload = createShellcodePayload(shellcode_byte_size, return_offset); - // Backup original data at injection address + // Backup original data at attach address std::vector backup_memory_data(shellcode_byte_size); - if (!process_tracer_.readMemory(injection_address, backup_memory_data.data(), + if (!process_tracer_.readMemory(attach_address, backup_memory_data.data(), shellcode_byte_size)) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::READ_TARGET_MEMORY_FAILED; } // Deploy shellcode - if (!process_tracer_.writeMemory(injection_address, shellcode_payload.data(), + if (!process_tracer_.writeMemory(attach_address, shellcode_payload.data(), shellcode_payload.size())) { // Restore state and detach on failure - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::WRITE_SHELLCODE_TO_TARGET_MEMORY_FAILED; } - // Now that the new code is in place, let the target run our injected code. + // Now that the new code is in place, let the target run our loaded code. if (!process_tracer_.continueExecution()) { // Restore state and detach on failure - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::ERROR_IN_EXECUTE_MALLOC; } // At this point, the target should have run malloc(). Check its return value. REG_TYPE malloc_registers_state; if (!process_tracer_.getRegisters(&malloc_registers_state)) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::GET_MALLOC_REGISTERS_FAILED; } unsigned long long target_buffer_address = malloc_registers_state.rax; if (target_buffer_address == 0) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::MALLOC_RETURN_ZERO; } @@ -385,73 +380,68 @@ ExitCode LibraryInjector::orchestrateInjectionSequence( if (!process_tracer_.writeMemory(target_buffer_address, library_file_path_.c_str(), library_path_string_length)) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::WRITE_LIBRARY_STR_TO_TARGET_MEMORY_FAILED; } // Continue the target's execution to call __libc_dlopen_mode. if (!process_tracer_.continueExecution()) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::ERROR_IN_EXECUTE_DLOPEN; } // Check the registers after calling dlopen. REG_TYPE dlopen_registers_state; if (!process_tracer_.getRegisters(&dlopen_registers_state)) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::GET_DLOPEN_REGISTERS_FAILED; } unsigned long long library_base_address = dlopen_registers_state.rax; if (library_base_address == 0) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::DLOPEN_RETURN_ZERO; } // As a courtesy, free the buffer that we allocated inside the target process. if (!process_tracer_.continueExecution()) { - process_tracer_.recoverInjection(injection_address, - backup_memory_data.data(), - shellcode_byte_size, initial_registers); + process_tracer_.recoverAttach(attach_address, backup_memory_data.data(), + shellcode_byte_size, initial_registers); return ExitCode::ERROR_IN_EXECUTE_FREE; } - // Confirm injection success and restore state - return confirmInjectionSuccess(injection_address, backup_memory_data, - shellcode_byte_size, initial_registers); + // Confirm attach success and restore state + return confirmAttachSuccess(attach_address, backup_memory_data, + shellcode_byte_size, initial_registers); } /** - * @brief Confirm the injection success and restore the original process state + * @brief Confirm the attach success and restore the original process state * * This function: - * 1. Restores the original process state using ProcessTracer's recoverInjection + * 1. Restores the original process state using ProcessTracer's recoverAttach * 2. Checks if the library was successfully loaded * - * @param injection_memory_location Address where the shellcode was injected + * @param attach_memory_location Address where the shellcode was loaded * @param backup_memory_data Vector containing the original data at the - * injection address + * attach address * @param shellcode_byte_size Size of the shellcode * @param original_register_state Pointer to the original register state - * @return true if the library was successfully loaded, false otherwise + * @return ExitCode indicating success or failure */ -ExitCode LibraryInjector::confirmInjectionSuccess( - long injection_memory_location, const std::vector &backup_memory_data, +ExitCode AttachAgent::confirmAttachSuccess( + long attach_memory_location, const std::vector &backup_memory_data, size_t shellcode_byte_size, REG_TYPE *original_register_state) { - // Restore the original state using ProcessTracer's recoverInjection for + // Restore the original state using ProcessTracer's recoverAttach for // consistency Create a temporary copy of the data for restoration std::vector temp_data(backup_memory_data); - if (!process_tracer_.recoverInjection(injection_memory_location, - temp_data.data(), shellcode_byte_size, - original_register_state)) { + if (!process_tracer_.recoverAttach(attach_memory_location, temp_data.data(), + shellcode_byte_size, + original_register_state)) { return ExitCode::ERROR_IN_EXECUTE_RECOVER_INJECTION; } @@ -481,7 +471,7 @@ ExitCode LibraryInjector::confirmInjectionSuccess( * * @param file_path Reference to the path string to modify */ -void LibraryInjector::getParentDirectoryPath(std::string &file_path) { +void AttachAgent::getParentDirectoryPath(std::string &file_path) { size_t last_slash_position = file_path.rfind('/'); if (last_slash_position != std::string::npos) { file_path.erase( diff --git a/csrc/attach/AttachAgent.h b/csrc/attach/AttachAgent.h new file mode 100644 index 0000000..b7ad2b3 --- /dev/null +++ b/csrc/attach/AttachAgent.h @@ -0,0 +1,72 @@ +#ifndef ATTACH_AGENT_H +#define ATTACH_AGENT_H + +#include "ExitCode.h" +#include "ProcessTracer.h" +#include "ProcessUtils.h" +#include +#include +#include +#include + +/** + * @brief Manages the attaching of profiler agent into target processes + * + * This class implements a comprehensive solution for attaching profiler agent + * into target processes using advanced ptrace-based techniques. It handles + * all aspects of the attach process including preparation, execution, and + * verification. + */ +class AttachAgent { +public: + /** + * @brief Constructs an AttachAgent instance + * @param target_process_id PID of the process to attach the agent into + * @param shared_library_file_path File path of the shared library to load + * @param debug_mode Enable debug logging + */ + AttachAgent(pid_t target_process_id, + const std::string &shared_library_file_path, + bool debug_mode = false); + + /** + * @brief Cleans up resources used by the AttachAgent + */ + ~AttachAgent(); + + /** + * @brief Performs the complete agent attach process + * @return ExitCode indicating success or failure + */ + ExitCode performAttach(); + +private: + // Core instance attributes + pid_t target_process_id_; + std::string library_file_path_; + ProcessTracer process_tracer_; + + // Attach workflow methods + ExitCode initializeAttachEnvironment(long &code_attach_address, + REG_TYPE *original_registers, + REG_TYPE *working_registers); + ExitCode orchestrateAttachSequence(long attach_address, + long malloc_function_address, + long free_function_address, + long dlopen_function_address, + int library_path_string_length, + REG_TYPE *initial_registers); + ExitCode confirmAttachSuccess(long attach_memory_location, + const std::vector &backup_memory_data, + size_t shellcode_byte_size, + REG_TYPE *original_register_state); + + // Shellcode generation methods + std::vector createShellcodePayload(size_t &payload_size, + intptr_t &return_instruction_offset); + + // Path manipulation utilities + void getParentDirectoryPath(std::string &file_path); +}; + +#endif // ATTACH_AGENT_H diff --git a/csrc/inject/ExitCode.h b/csrc/attach/ExitCode.h similarity index 100% rename from csrc/inject/ExitCode.h rename to csrc/attach/ExitCode.h diff --git a/csrc/inject/ProcessTracer.cpp b/csrc/attach/ProcessTracer.cpp similarity index 93% rename from csrc/inject/ProcessTracer.cpp rename to csrc/attach/ProcessTracer.cpp index 7ccdb50..5aec348 100644 --- a/csrc/inject/ProcessTracer.cpp +++ b/csrc/attach/ProcessTracer.cpp @@ -272,28 +272,27 @@ bool ProcessTracer::verifySignalStatus() { } /** - * @brief Restore the process state after a failed injection attempt + * @brief Restore the process state after a failed attach attempt * * This function performs the three required steps to restore the process state: - * 1. Write the original memory data back to the injection address + * 1. Write the original memory data back to the attach address * 2. Restore the original register state * 3. Detach from the process * - * @param injection_address Address where the shellcode was injected - * @param backup_data Pointer to the original data at the injection address + * @param attach_address Address where the shellcode was loaded + * @param backup_data Pointer to the original data at the attach address * @param data_length Length of the backup data * @param registers Pointer to the original register state * @return true if restoration was successful, false otherwise */ -bool ProcessTracer::recoverInjection(long injection_address, - const void *backup_data, - size_t data_length, REG_TYPE *registers) { - // Step 1: Write the original memory data back to the injection address - if (!writeMemory(injection_address, backup_data, data_length)) { +bool ProcessTracer::recoverAttach(long attach_address, const void *backup_data, + size_t data_length, REG_TYPE *registers) { + // Step 1: Write the original memory data back to the attach address + if (!writeMemory(attach_address, backup_data, data_length)) { if (debug_mode_) { std::cerr << "[ERROR] PyFlightProfiler: Failed to recover original " "memory data at address 0x" - << std::hex << injection_address << std::dec << std::endl; + << std::hex << attach_address << std::dec << std::endl; } return false; } diff --git a/csrc/inject/ProcessTracer.h b/csrc/attach/ProcessTracer.h similarity index 91% rename from csrc/inject/ProcessTracer.h rename to csrc/attach/ProcessTracer.h index 91d83e6..2092769 100644 --- a/csrc/inject/ProcessTracer.h +++ b/csrc/attach/ProcessTracer.h @@ -112,15 +112,15 @@ class ProcessTracer { // Failure recovery /** - * @brief Restore the process state after a failed injection attempt - * @param injection_address Address where the shellcode was injected - * @param backup_data Pointer to the original data at the injection address + * @brief Restore the process state after a failed attach attempt + * @param attach_address Address where the shellcode was loaded + * @param backup_data Pointer to the original data at the attach address * @param data_length Length of the backup data * @param registers Pointer to the original register state * @return true if restoration was successful, false otherwise */ - bool recoverInjection(long injection_address, const void *backup_data, - size_t data_length, REG_TYPE *registers); + bool recoverAttach(long attach_address, const void *backup_data, + size_t data_length, REG_TYPE *registers); // Accessor for debug mode /** diff --git a/csrc/inject/ProcessUtils.cpp b/csrc/attach/ProcessUtils.cpp similarity index 100% rename from csrc/inject/ProcessUtils.cpp rename to csrc/attach/ProcessUtils.cpp diff --git a/csrc/inject/ProcessUtils.h b/csrc/attach/ProcessUtils.h similarity index 100% rename from csrc/inject/ProcessUtils.h rename to csrc/attach/ProcessUtils.h diff --git a/csrc/inject/inject.cpp b/csrc/attach/attach.cpp similarity index 77% rename from csrc/inject/inject.cpp rename to csrc/attach/attach.cpp index dc39692..f532c25 100644 --- a/csrc/inject/inject.cpp +++ b/csrc/attach/attach.cpp @@ -1,4 +1,4 @@ -#include "LibraryInjector.h" +#include "AttachAgent.h" #include "ProcessUtils.h" #include #include @@ -26,22 +26,22 @@ void extractParentDirectoryFromPath(std::string &file_system_path) { } /** - * @brief Main entry point for the library injection utility + * @brief Main entry point for the profiler attach utility * - * This program injects the flight_profiler_agent.so library into a target - * process using advanced ptrace-based injection techniques. + * This program attaches the flight_profiler_agent.so library into a target + * process using advanced ptrace-based techniques. * - * Usage: ./inject + * Usage: ./attach * * @param argument_count Number of command line arguments * @param argument_values Array of command line arguments - * @return 0 on successful injection, 1 on failure + * @return 0 on successful attach, 1 on failure */ int main(int argument_count, char **argument_values) { // Validate command line arguments if (argument_count < 2) { - std::cout << "Invalid inject command without target process identifier " - "provided, USAGE: ./inject process_id!" + std::cout << "Invalid attach command without target process identifier " + "provided, USAGE: ./attach process_id!" << std::endl; return 1; } @@ -86,9 +86,8 @@ int main(int argument_count, char **argument_values) { std::string library_file_path(library_file_path_cstring); free(library_file_path_cstring); - // Create and execute the injector - LibraryInjector library_injector(target_process_id, library_file_path, - debug_mode); + // Create and execute the attach agent + AttachAgent attach_agent(target_process_id, library_file_path, debug_mode); - return static_cast(library_injector.performInjection()); + return static_cast(attach_agent.performAttach()); } diff --git a/csrc/code_inject.h b/csrc/code_inject.h deleted file mode 100644 index 563e76d..0000000 --- a/csrc/code_inject.h +++ /dev/null @@ -1,16 +0,0 @@ - - -#ifndef __CODE_INJECT_H__ -#define __CODE_INJECT_H__ - -#ifdef __cplusplus -extern "C" { -#endif - -int inject(char *fn, int p, unsigned long nm_symbol_offset); - -#ifdef __cplusplus -} -#endif - -#endif diff --git a/csrc/frida_profiler.cpp b/csrc/frida_profiler.cpp index 3f2a10d..ed90554 100644 --- a/csrc/frida_profiler.cpp +++ b/csrc/frida_profiler.cpp @@ -16,13 +16,12 @@ int init_frida_gum() { pthread_mutex_lock(&mutex); if (inited != 0) { pthread_mutex_unlock(&mutex); - fprintf(stderr, "[*] frida gum already inited\n"); return 0; } gum_init_embedded(); inited = 1; pthread_mutex_unlock(&mutex); - g_print("[*] init frida gum successfully\n"); + fprintf(stdout, "[PyFlightProfiler] Native profiler initialized.\n"); return 0; } @@ -30,13 +29,12 @@ int deinit_frida_gum() { pthread_mutex_lock(&mutex); if (inited != 1) { // pthread_mutex_unlock(&mutex); - fprintf(stderr, "[*] frida gum not inited\n"); return 0; } gum_deinit_embedded(); inited = 0; pthread_mutex_unlock(&mutex); - g_print("[*] deinit frida gum successfully\n"); + fprintf(stdout, "[PyFlightProfiler] Native profiler deinitialized.\n"); return 0; } diff --git a/csrc/inject/LibraryInjector.h b/csrc/inject/LibraryInjector.h deleted file mode 100644 index f13c585..0000000 --- a/csrc/inject/LibraryInjector.h +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef LIBRARY_INJECTOR_H -#define LIBRARY_INJECTOR_H - -#include "ExitCode.h" -#include "ProcessTracer.h" -#include "ProcessUtils.h" -#include -#include -#include -#include - -/** - * @brief Manages the injection of shared libraries into target processes - * - * This class implements a comprehensive solution for injecting shared libraries - * into target processes using advanced ptrace-based techniques. It handles - * all aspects of the injection process including preparation, execution, and - * verification. - */ -class LibraryInjector { -public: - /** - * @brief Constructs a LibraryInjector instance - * @param target_process_id PID of the process to inject the library into - * @param shared_library_file_path File path of the shared library to inject - * @param debug_mode Enable debug logging - */ - LibraryInjector(pid_t target_process_id, - const std::string &shared_library_file_path, - bool debug_mode = false); - - /** - * @brief Cleans up resources used by the LibraryInjector - */ - ~LibraryInjector(); - - /** - * @brief Performs the complete library injection process - * @return true if the injection completed successfully, false otherwise - */ - ExitCode performInjection(); - -private: - // Core instance attributes - pid_t target_process_id_; - std::string library_file_path_; - ProcessTracer process_tracer_; - - // Injection workflow methods - ExitCode initializeInjectionEnvironment(long &code_injection_address, - REG_TYPE *original_registers, - REG_TYPE *working_registers); - ExitCode orchestrateInjectionSequence(long injection_address, - long malloc_function_address, - long free_function_address, - long dlopen_function_address, - int library_path_string_length, - REG_TYPE *initial_registers); - ExitCode confirmInjectionSuccess(long injection_memory_location, - const std::vector &backup_memory_data, - size_t shellcode_byte_size, - REG_TYPE *original_register_state); - - // Shellcode generation methods - std::vector createShellcodePayload(size_t &payload_size, - intptr_t &return_instruction_offset); - - // Path manipulation utilities - void getParentDirectoryPath(std::string &file_path); -}; - -#endif // LIBRARY_INJECTOR_H diff --git a/csrc/code_inject.cpp b/csrc/profiler_attach.cpp similarity index 85% rename from csrc/code_inject.cpp rename to csrc/profiler_attach.cpp index d294066..6970296 100644 --- a/csrc/code_inject.cpp +++ b/csrc/profiler_attach.cpp @@ -1,4 +1,4 @@ -#include "code_inject.h" +#include "profiler_attach.h" #include "Python.h" #include "frida_profiler.h" #include "py_gil_intercept.h" @@ -11,7 +11,7 @@ #include static char filename[FILENAME_MAX] = ""; static char take_gil_literal[9] = "take_gil"; -static int py_injected = 0; +static int py_attached = 0; static int port; static pthread_mutex_t mutex; @@ -99,7 +99,7 @@ static void boot_entry(void *boot_raw) { if (tstate == NULL) { PyMem_DEL(boot_raw); fprintf(stderr, - "pyFlightProfiler: Not enough memory to create thread state.\n"); + "[PyFlightProfiler] Not enough memory to create thread state.\n"); return; } @@ -107,14 +107,14 @@ static void boot_entry(void *boot_raw) { // here take gil lock, and set current PyThreadState PyEval_AcquireThread(tstate); - fprintf(stdout, "pyFlightProfiler: CodeInject Executing %s.\n", filename); + fprintf(stdout, "[PyFlightProfiler] Loading agent: %s\n", filename); PyObject *res = exec_python_entrance(); if (res == NULL) { if (PyErr_ExceptionMatches(PyExc_SystemExit)) { /* SystemExit is ignored silently */ PyErr_Clear(); } else { - fprintf(stderr, "pyFlightProfiler: Unhandled exception in thread.\n"); + fprintf(stderr, "[PyFlightProfiler] Unhandled exception in thread.\n"); PyErr_PrintEx(0); // clear state, we don't want to crash the other process PyErr_Clear(); @@ -123,7 +123,7 @@ static void boot_entry(void *boot_raw) { Py_DECREF(res); } - fprintf(stdout, "pyFlightProfiler: Thread finished execution.\n"); + fprintf(stdout, "[PyFlightProfiler] Agent initialization complete.\n"); PyMem_RawFree(boot_raw); // clear tstat data @@ -142,7 +142,7 @@ static int start_thread() { // PyMem_NEW or PyMem_Malloc not work in python 3.12 boot = (struct bootstate *)PyMem_RawMalloc(sizeof(struct bootstate)); if (boot == NULL) { - fprintf(stderr, "pyFlightProfiler: alloc memory for bootstate failed\n"); + fprintf(stderr, "[PyFlightProfiler] alloc memory for bootstate failed\n"); return 1; } @@ -154,7 +154,7 @@ static int start_thread() { PyGILState_STATE old_gil_state = PyGILState_Ensure(); boot->interp = PyThreadState_Get()->interp; - // start a background thread to inject code, gdb will quit quickly + // start a background thread to attach code, debugger will quit quickly ident = PyThread_start_new_thread(boot_entry, (void *)boot); int ret = 0; @@ -174,12 +174,12 @@ static int start_thread() { extern "C" { #endif -int inject(char *fn, int p, unsigned long nm_symbol_offset) { +int profiler_attach(char *fn, int p, unsigned long nm_symbol_offset) { int ret_port = p; int ret; pthread_mutex_lock(&mutex); - if (py_injected != 0) { - fprintf(stderr, "code already injected"); + if (py_attached != 0) { + fprintf(stderr, "profiler already attached"); ret_port = port; } else { strcpy(filename, fn); @@ -189,7 +189,7 @@ int inject(char *fn, int p, unsigned long nm_symbol_offset) { pthread_mutex_unlock(&mutex); return -1; } - py_injected = 1; + py_attached = 1; } pthread_mutex_unlock(&mutex); // init offset from proc addr to addr get from system nm command @@ -202,7 +202,7 @@ int inject(char *fn, int p, unsigned long nm_symbol_offset) { return ret_port; } -int inject_init_frida_gum(unsigned long nm_symbol_offset) { +int init_native_profiler(unsigned long nm_symbol_offset) { // used for CPython >= 3.14 // init offset from proc addr to addr get from system nm command set_nm_symbol_offset(nm_symbol_offset); @@ -236,7 +236,7 @@ const void *get_so_path() { return dl_info.dli_fname; } -void inject_inner() { +void do_attach() { pthread_mutex_init(&mutex, NULL); const char *so_path = (const char *)get_so_path(); @@ -251,7 +251,7 @@ void inject_inner() { get_parent_directory(so_path_modify); char params_path[PATH_MAX]; // Make sure the buffer is large enough - snprintf(params_path, sizeof(params_path), "%s/inject_params.data", + snprintf(params_path, sizeof(params_path), "%s/attach_params.data", so_path_modify); FILE *file; @@ -262,7 +262,7 @@ void inject_inner() { file = fopen(params_path, "r"); if (file == NULL) { - perror("Unable to open input_prams.data"); + perror("Unable to open attach_params.data"); return; } if (fgets(line, sizeof(line), file) != NULL) { @@ -285,11 +285,11 @@ void inject_inner() { base_addr = strtoul(token, NULL, 10); } } else { - fprintf(stderr, "Error reading input_params.data!\n"); + fprintf(stderr, "Error reading attach_params.data!\n"); return; } fclose(file); - inject(py_code, port, base_addr); + profiler_attach(py_code, port, base_addr); } #ifdef __cplusplus @@ -297,10 +297,10 @@ void inject_inner() { #endif /* - * This function is automatically called when the library is injected + * This function is automatically called when the library is loaded * into a process. */ -__attribute__((constructor)) void code_inject_init() { +__attribute__((constructor)) void profiler_attach_init() { #if !defined(__APPLE__) if (Py_IsInitialized()) { if (Py_GetVersion() != NULL) { @@ -308,13 +308,13 @@ __attribute__((constructor)) void code_inject_init() { int major, minor; if (sscanf(version, "%d.%d", &major, &minor) == 2) { if (major == 3 && minor < 14) { - inject_inner(); + do_attach(); } } } } else { #if PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 14 - inject_inner(); + do_attach(); #endif } #endif diff --git a/csrc/profiler_attach.h b/csrc/profiler_attach.h new file mode 100644 index 0000000..8873802 --- /dev/null +++ b/csrc/profiler_attach.h @@ -0,0 +1,16 @@ + + +#ifndef __PROFILER_ATTACH_H__ +#define __PROFILER_ATTACH_H__ + +#ifdef __cplusplus +extern "C" { +#endif + +int profiler_attach(char *fn, int p, unsigned long nm_symbol_offset); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/csrc/py_gil_stat.cpp b/csrc/py_gil_stat.cpp index a42f87a..02a2db4 100644 --- a/csrc/py_gil_stat.cpp +++ b/csrc/py_gil_stat.cpp @@ -340,11 +340,11 @@ void PyGilStat::boot_entry(void *boot_raw) { if (tstate == NULL) { PyMem_DEL(boot_raw); fprintf(stderr, - "pyFlightProfiler: Not enough memory to create thread state.\n"); + "[PyFlightProfiler] Not enough memory to create thread state.\n"); return; } - fprintf(stdout, "pyFlightProfiler: Gil Stat Thread start Executing.\n"); + fprintf(stdout, "[PyFlightProfiler] Gil Stat Thread start Executing.\n"); int stat_interval = stat->config->stat_interval * 1000; int sleep_interval = 500; @@ -382,7 +382,7 @@ void PyGilStat::boot_entry(void *boot_raw) { select(0, (fd_set *)0, (fd_set *)0, (fd_set *)0, &t); } - fprintf(stdout, "pyFlightProfiler: Gil Stat Thread finished execution.\n"); + fprintf(stdout, "[PyFlightProfiler] Gil Stat Thread finished execution.\n"); PyMem_RawFree(boot_raw); @@ -406,7 +406,7 @@ void PyGilStat::start_python_stat_thread() { boot = (struct bootstate *)PyMem_RawMalloc(sizeof(struct bootstate)); if (boot == NULL) { fprintf(stderr, - "pyFlightProfiler: alloc memory for gil stat bootstate failed\n"); + "[PyFlightProfiler] alloc memory for gil stat bootstate failed\n"); PyErr_PrintEx(0); return; } diff --git a/flight_profiler/client.py b/flight_profiler/client.py index 76cb325..ffa9aa1 100644 --- a/flight_profiler/client.py +++ b/flight_profiler/client.py @@ -4,6 +4,7 @@ import os import platform import re +import shutil import signal import socket import sys @@ -30,13 +31,17 @@ ) from flight_profiler.utils.env_util import is_linux, is_mac, py_higher_than_314 from flight_profiler.utils.render_util import ( + BANNER_COLOR_CYAN, + BOX_HORIZONTAL, + COLOR_BRIGHT_GREEN, COLOR_END, + COLOR_FAINT, COLOR_GREEN, COLOR_ORANGE, COLOR_RED, COLOR_WHITE_255, - build_colorful_banners, - build_title_hints, + build_prompt_separator, + build_welcome_box, ) from flight_profiler.utils.shell_util import execute_shell, get_py_bin_path @@ -47,6 +52,349 @@ except ImportError: READLINE_AVAILABLE = False +# Check termios/tty availability for advanced terminal input (Unix only) +try: + import termios + import tty + TERMIOS_AVAILABLE = True +except ImportError: + TERMIOS_AVAILABLE = False + + +def get_cursor_position() -> int: + """ + Get current cursor row position in terminal (1-based). + Returns -1 if unable to detect. + """ + if not TERMIOS_AVAILABLE: + return -1 + try: + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + try: + tty.setraw(fd) + sys.stdout.write('\033[6n') + sys.stdout.flush() + response = '' + while True: + ch = sys.stdin.read(1) + response += ch + if ch == 'R': + break + # Response format: \033[row;colR + match = re.search(r'\[(\d+);(\d+)R', response) + if match: + return int(match.group(1)) + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + except Exception: + pass + return -1 + + +def ensure_space_from_bottom(min_lines: int = 3) -> None: + """ + Ensure there's enough space from the bottom of terminal. + If cursor is too close to bottom, scroll up by printing newlines. + """ + try: + terminal_height = shutil.get_terminal_size().lines + cursor_row = get_cursor_position() + if cursor_row > 0: + lines_from_bottom = terminal_height - cursor_row + if lines_from_bottom < min_lines: + # Need to scroll up + scroll_lines = min_lines - lines_from_bottom + print('\n' * scroll_lines, end='') + # Move cursor back up + sys.stdout.write(f'\033[{scroll_lines}A') + sys.stdout.flush() + except Exception: + pass + + +def read_input_with_box(prompt: str, prompt_gray: str, show_placeholder: bool = False) -> str: + """ + Read single-line input with a box frame (top and bottom separators). + The input area appears between two horizontal lines. + Enter submits, Ctrl-D exits. + After submission, clears the box and changes prompt to gray. + + Falls back to standard input() if termios is not available. + """ + terminal_width = shutil.get_terminal_size().columns + separator = f"{COLOR_FAINT}{BOX_HORIZONTAL * terminal_width}{COLOR_END}" + placeholder = "help" + + # Fallback for systems without termios (e.g., Windows) + if not TERMIOS_AVAILABLE: + print(separator) + result = input(prompt).strip() + print(separator) + return result + + # Print the box frame: top line, input line placeholder, bottom line + print(separator) # Top separator + sys.stdout.write(prompt) # Prompt + # Show placeholder if requested + if show_placeholder: + sys.stdout.write(f'{COLOR_FAINT}{placeholder}{COLOR_END}') + sys.stdout.write('\n') # Move to next line + print(separator) # Bottom separator + + # Move cursor back up to the input line (2 lines up, then to prompt position) + sys.stdout.write('\033[2A') # Move up 2 lines + prompt_len = 2 # ❯ + space (❯ is 1 width char) + sys.stdout.write(f'\033[{prompt_len + 1}G') # Move to position after prompt + sys.stdout.flush() + + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + + line = '' + cursor_pos = 0 + ctrl_d_pressed = False # Track if Ctrl-D was pressed once + ctrl_c_pressed = False # Track if Ctrl-C was pressed once + placeholder_visible = show_placeholder # Track if placeholder is currently shown + + # History navigation + history_index = -1 # -1 means current input, 0 is most recent history + saved_line = '' # Save current input when navigating history + + def get_history_length(): + """Get the number of history entries.""" + if READLINE_AVAILABLE: + return readline.get_current_history_length() + return 0 + + def get_history_item(index): + """Get history item by index (1-based in readline).""" + if READLINE_AVAILABLE and index > 0: + return readline.get_history_item(index) + return None + + def replace_line(new_text): + """Replace current line with new text and update display.""" + nonlocal line, cursor_pos + # Clear current line content + sys.stdout.write('\r') + sys.stdout.write(prompt) + sys.stdout.write(' ' * len(line)) + sys.stdout.write('\r') + sys.stdout.write(prompt) + # Write new content + sys.stdout.write(new_text) + sys.stdout.flush() + line = new_text + cursor_pos = len(line) + + def cleanup_box_and_show_result(input_text: str): + """Clear the box frame and show the command with gray prompt.""" + # Current cursor is on the input line (line 2) + # Line 1: top separator + # Line 2: input line (cursor here) + # Line 3: bottom separator + + # Move to start of line + sys.stdout.write('\r') + # Move up to top separator (line 1) + sys.stdout.write('\033[1A') + # Clear top separator line + sys.stdout.write('\033[2K') + # Print gray prompt + input text (replaces top separator) + sys.stdout.write(f'{prompt_gray}{input_text}') + # Move down to line 2 (old input line) + sys.stdout.write('\n\033[2K') # Clear old input line + # Move down to line 3 (bottom separator), command output starts here + sys.stdout.write('\n\033[2K') # Clear bottom separator + sys.stdout.flush() + + try: + # Set terminal to cbreak mode but also disable ISIG to capture Ctrl-C as character + new_settings = termios.tcgetattr(fd) + new_settings[3] = new_settings[3] & ~termios.ECHO & ~termios.ICANON & ~termios.ISIG # lflags + new_settings[6][termios.VMIN] = 1 + new_settings[6][termios.VTIME] = 0 + termios.tcsetattr(fd, termios.TCSADRAIN, new_settings) + + while True: + ch = sys.stdin.read(1) + + if ch == '\x04': # Ctrl-D + if not line.strip(): + if ctrl_d_pressed: + # Second Ctrl-D - exit silently, clear hint first + sys.stdout.write('\033[1B') # Move to bottom separator line + sys.stdout.write('\n\033[2K') # Move down and clear hint line + sys.stdout.write('\n') # Add extra newline at end + sys.stdout.flush() + raise EOFError() + else: + # First Ctrl-D - show hint below bottom separator + ctrl_d_pressed = True + sys.stdout.write('\033[1B') # Move down to bottom separator line + sys.stdout.write('\n') # Move to line below + sys.stdout.write(f'{COLOR_FAINT}Press Ctrl-D again to exit{COLOR_END}') + sys.stdout.write('\033[2A') # Move back up 2 lines to input line + sys.stdout.write(f'\033[{prompt_len + cursor_pos + 1}G') # Restore cursor position + sys.stdout.flush() + else: + # Submit with Ctrl-D if there's content + cleanup_box_and_show_result(line) + return line.strip() + + elif ch == '\x03': # Ctrl-C + if ctrl_c_pressed: + # Second Ctrl-C - exit silently, clear hint but keep bottom separator + sys.stdout.write('\r\033[2K') # Clear current line first + sys.stdout.write('\033[1B') # Move to bottom separator line + sys.stdout.write('\n\033[2K') # Move down and clear hint line + sys.stdout.write('\n') # Add extra newline at end + sys.stdout.flush() + raise KeyboardInterrupt + else: + # First Ctrl-C - clear input and show hint + ctrl_c_pressed = True + # Clear current line and rewrite prompt (removes any ^C echo) + sys.stdout.write('\r\033[2K') # Clear current line + sys.stdout.write(prompt) # Rewrite prompt + line = '' + cursor_pos = 0 + # Clear any previous Ctrl-D hint + if ctrl_d_pressed: + ctrl_d_pressed = False + # Show Ctrl-C hint below bottom separator + sys.stdout.write('\033[1B') # Move down to bottom separator line + sys.stdout.write('\n') # Move to line below (don't clear separator!) + sys.stdout.write(f'{COLOR_FAINT}Press Ctrl-C again to exit{COLOR_END}') + sys.stdout.write('\033[2A') # Move back up 2 lines to input line + sys.stdout.write(f'\033[{prompt_len + 1}G') # Move to position after prompt + sys.stdout.flush() + + elif ch == '\n' or ch == '\r': # Enter - submit only if has input + if line.strip(): + cleanup_box_and_show_result(line) + return line.strip() + # Empty input - do nothing, stay in place + + elif ch == '\x7f' or ch == '\x08': # Backspace + if cursor_pos > 0: + line = line[:cursor_pos-1] + line[cursor_pos:] + cursor_pos -= 1 + # Move back, clear to end of line, reprint rest + sys.stdout.write('\b') + rest = line[cursor_pos:] + sys.stdout.write(rest + ' ') + sys.stdout.write('\b' * (len(rest) + 1)) + sys.stdout.flush() + + elif ch == '\033': # Escape sequence (arrow keys) + seq1 = sys.stdin.read(1) + if seq1 == '[': + seq2 = sys.stdin.read(1) + if seq2 == 'D': # Left arrow + if cursor_pos > 0: + cursor_pos -= 1 + sys.stdout.write('\033[D') + sys.stdout.flush() + elif seq2 == 'C': # Right arrow + if cursor_pos < len(line): + cursor_pos += 1 + sys.stdout.write('\033[C') + sys.stdout.flush() + elif seq2 == 'A': # Up arrow - previous history + history_len = get_history_length() + if history_len > 0: + # Clear placeholder if visible + if placeholder_visible: + placeholder_visible = False + sys.stdout.write('\033[2K') + sys.stdout.write('\r') + sys.stdout.write(prompt) + sys.stdout.flush() + # Save current line when first navigating + if history_index == -1: + saved_line = line + # Move to older history + if history_index < history_len - 1: + history_index += 1 + hist_item = get_history_item(history_len - history_index) + if hist_item: + replace_line(hist_item) + elif seq2 == 'B': # Down arrow - next history + if history_index > -1: + # Clear placeholder if visible + if placeholder_visible: + placeholder_visible = False + sys.stdout.write('\033[2K') + sys.stdout.write('\r') + sys.stdout.write(prompt) + sys.stdout.flush() + history_index -= 1 + if history_index == -1: + # Back to current input + replace_line(saved_line) + else: + history_len = get_history_length() + hist_item = get_history_item(history_len - history_index) + if hist_item: + replace_line(hist_item) + + elif ch == '\t': # Tab - command completion + words = line.strip().split() + # Only complete first command when no space after it + if len(words) <= 1 and (len(line) == 0 or not line.endswith(' ')): + prefix = words[0] if words else '' + matches = [name for name in HELP_COMMANDS_NAMES if name.startswith(prefix)] + if len(matches) == 1: + # Single match - auto complete + completion = matches[0] + ' ' + # Clear current input and show completed text + sys.stdout.write('\b' * cursor_pos) + sys.stdout.write(' ' * len(line)) + sys.stdout.write('\b' * len(line)) + sys.stdout.write(completion) + sys.stdout.flush() + line = completion + cursor_pos = len(line) + elif len(matches) > 1: + # Multiple matches - show options below + sys.stdout.write('\n') + sys.stdout.write(f"{COLOR_FAINT} {' '.join(matches)}{COLOR_END}") + sys.stdout.write(f'\033[1A') # Move up 1 line + sys.stdout.write(f'\033[{prompt_len + cursor_pos + 1}G') # Restore cursor position + sys.stdout.flush() + + elif ch >= ' ' and ch <= '~': # Printable character + # Clear placeholder on first input + if placeholder_visible: + placeholder_visible = False + # Clear placeholder text + sys.stdout.write('\033[2K') # Clear current line + sys.stdout.write('\r') + sys.stdout.write(prompt) # Rewrite prompt + sys.stdout.flush() + # Reset Ctrl-D/Ctrl-C state and clear hint if shown + if ctrl_d_pressed or ctrl_c_pressed: + ctrl_d_pressed = False + ctrl_c_pressed = False + # Clear the hint below the box + sys.stdout.write('\033[1B') # Move to bottom separator + sys.stdout.write('\n\033[2K') # Move down and clear hint line + sys.stdout.write('\033[2A') # Move back up to input line + sys.stdout.write(f'\033[{prompt_len + cursor_pos + 1}G') # Restore cursor position + line = line[:cursor_pos] + ch + line[cursor_pos:] + cursor_pos += 1 + sys.stdout.write(ch) + rest = line[cursor_pos:] + if rest: + sys.stdout.write(rest) + sys.stdout.write('\b' * len(rest)) + sys.stdout.flush() + + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + class ProfilerCli(object): def __init__(self, port: int, @@ -61,29 +409,44 @@ def __init__(self, port: int, self.history_file = os.path.join(output_dir, "cli_history") set_history_file_path(self.history_file) self.current_plugin = None + self.first_input = True # Track if this is the first command input def run(self): - build_colorful_banners() - build_title_hints([ - ("pid", str(self.server_pid)), - ("py_executable", self.target_executable) - ]) + build_welcome_box(str(self.server_pid), self.target_executable) while True: try: - prompt = f"[cmd@{self.server_pid}]$ " - cmd = input(prompt).strip() + # Ensure there's space from terminal bottom (at least 5 lines for input box) + ensure_space_from_bottom(5) + + # White/bright prompt for active input, gray for history + prompt_active = f"{COLOR_WHITE_255}❯{COLOR_END} " + prompt_gray = f"{COLOR_FAINT}❯{COLOR_END} " + + # Read input with box frame (Enter to submit, Ctrl-D to exit) + # Show placeholder hint only on first input + cmd = read_input_with_box(prompt_active, prompt_gray, show_placeholder=self.first_input) + if len(cmd) == 0: - print("", end="") continue + + # After first successful command, don't show placeholder anymore + self.first_input = False + + # Add to history if readline is available + if READLINE_AVAILABLE: + readline.add_history(cmd) + self.do_action(cmd) except EOFError: if READLINE_AVAILABLE: readline.write_history_file(self.history_file) - sys.exit("CTRL+D pressed. Exiting Profiler.") + sys.exit(0) except KeyboardInterrupt: - print("") - pass + # Second Ctrl-C in input - exit silently + if READLINE_AVAILABLE: + readline.write_history_file(self.history_file) + sys.exit(0) def check_need_help(self, cmd: str) -> bool: return " --help " in cmd or " -h " in cmd or cmd.endswith("-h") or cmd.endswith("--help") @@ -112,7 +475,7 @@ def do_action(self, cmd: str): except ModuleNotFoundError as e: print( f"{COLOR_RED} Unsupported command {parts[0]}, use {COLOR_END}{COLOR_ORANGE}help{COLOR_END}{COLOR_RED} " - f"to find available commands!{COLOR_END}" + f"to find available commands!{COLOR_END}\n" ) return self.current_plugin = module.get_instance(self.port, self.server_pid) @@ -129,10 +492,12 @@ def do_action(self, cmd: str): cmd[cmd.find(parts[0]) + len(parts[0]) :] ) except KeyboardInterrupt: + # Clear ^C from terminal and add newline + sys.stdout.write('\r\033[2K\n') + sys.stdout.flush() if self.current_plugin is not None: try: self.current_plugin.on_interrupted() - print() # create new line except Exception: show_error_info(traceback.format_exc()) except Exception: @@ -290,7 +655,7 @@ def get_base_addr(current_directory: str, server_pid: str, platform: str) -> int return base_addr -def do_inject_on_linux(free_port: int, server_pid: str, debug: bool = False) -> int: +def do_inject_on_linux(free_port: int, server_pid: str, debug: bool = False, diagnostic_messages: list = None) -> int: """ inject by ptrace under linux env returns target port if inject successfully, otherwise exit abnormally @@ -298,11 +663,11 @@ def do_inject_on_linux(free_port: int, server_pid: str, debug: bool = False) -> current_directory = os.path.dirname(os.path.abspath(__file__)) base_addr = get_base_addr(current_directory, server_pid, "linux") - code_inject_py: str = os.path.join(current_directory, "code_inject.py") - with open(os.path.join(current_directory, "lib/inject_params.data"), "w") as f: - f.write(f"{code_inject_py.strip()},{free_port},{base_addr}\n") + profiler_agent_py: str = os.path.join(current_directory, "profiler_agent.py") + with open(os.path.join(current_directory, "lib/attach_params.data"), "w") as f: + f.write(f"{profiler_agent_py.strip()},{free_port},{base_addr}\n") - shell_path = os.path.join(current_directory, "lib/inject") + shell_path = os.path.join(current_directory, "lib/attach") # Add debug flag to the command if enabled cmd_args = [str(shell_path), server_pid] if debug: @@ -317,18 +682,21 @@ def do_inject_on_linux(free_port: int, server_pid: str, debug: bool = False) -> text=True, ) exit_code = ps.wait() + if exit_code != 0 and not debug and diagnostic_messages: + for msg in diagnostic_messages: + print(msg) verify_exit_code(exit_code, server_pid) return free_port -def do_inject_on_mac(free_port: int, server_pid: str, debug: bool = False) -> int: +def do_inject_on_mac(free_port: int, server_pid: str, debug: bool = False, diagnostic_messages: list = None) -> int: """ inject by lldb under mac env returns target port if inject successfully, otherwise exit abnormally """ tmp_fd, tmp_file_path = tempfile.mkstemp() current_directory = os.path.dirname(os.path.abspath(__file__)) - shell_path = os.path.join(current_directory, "shell/code_inject.sh") + shell_path = os.path.join(current_directory, "shell/profiler_attach.sh") # Prepare command arguments cmd_args = [str(shell_path), str(os.getpid()), server_pid, tmp_file_path, str(free_port)] @@ -343,28 +711,29 @@ def do_inject_on_mac(free_port: int, server_pid: str, debug: bool = False) -> in bufsize=1, text=True, ) - if debug: - # Only print output in debug mode - while True: - output = ps.stdout.readline() - if output: - print(output, end="") - else: - break - ps.wait() + stdout_output, stderr_output = ps.communicate() + if debug and stdout_output: + print(stdout_output, end="") with open(tmp_file_path, "r") as temp: content = temp.read() if content is None or len(content) == 0: + # Print diagnostic info on failure (skip if already printed in debug mode) + if not debug and diagnostic_messages: + for msg in diagnostic_messages: + print(msg) + # Print process output on failure (contains error details) + if not debug and stdout_output: + print(stdout_output, end="" if stdout_output.endswith("\n") else "\n") print("PyFlightProfiler attach failed!") exit(1) return int(content) -def do_inject_with_sys_remote_exec(free_port: int, server_pid: str, debug: bool = False): +def do_inject_with_sys_remote_exec(free_port: int, server_pid: str, debug: bool = False, diagnostic_messages: list = None): current_directory = os.path.dirname(os.path.abspath(__file__)) - code_inject_py: str = os.path.join(current_directory, "code_inject.py") - inject_code_file_path: str = os.path.join(current_directory, f"code_inject_{server_pid}_{int(time.time())}.py") + profiler_agent_py: str = os.path.join(current_directory, "profiler_agent.py") + inject_code_file_path: str = os.path.join(current_directory, f"profiler_agent_{server_pid}_{int(time.time())}.py") shared_lib_suffix = "so" if is_linux() else "dylib" inject_agent_so_path: str = os.path.join(current_directory, "lib", f"flight_profiler_agent.{shared_lib_suffix}") @@ -373,7 +742,7 @@ def do_inject_with_sys_remote_exec(free_port: int, server_pid: str, debug: bool else: nm_symbol_offset = get_base_addr(current_directory, server_pid, "mac") - with open(code_inject_py, 'r', encoding='utf-8') as f: + with open(profiler_agent_py, 'r', encoding='utf-8') as f: content = f.read() modified_content = content.replace("${listen_port}", str(free_port)) modified_content = modified_content.replace("${current_file_abspath}", inject_code_file_path) @@ -385,21 +754,34 @@ def do_inject_with_sys_remote_exec(free_port: int, server_pid: str, debug: bool try: sys.remote_exec(int(server_pid), inject_code_file_path) except PermissionError as e: + # Print diagnostic info on failure (skip if already printed in debug mode) + if not debug and diagnostic_messages: + for msg in diagnostic_messages: + print(msg) show_error_info(f"\n[ERROR] Higher Permission required! This error id caused by {e}") show_normal_info(f"[{COLOR_GREEN}Solution{COLOR_END}{COLOR_WHITE_255}] Try run flight_profiler $pid as {COLOR_RED}root{COLOR_END}{COLOR_WHITE_255}!") exit(1) except: + # Print diagnostic info on failure (skip if already printed in debug mode) + if not debug and diagnostic_messages: + for msg in diagnostic_messages: + print(msg) logger.exception(f"Attach via sys.remote_exec failed!") exit(1) return free_port -def show_pre_attach_info(server_pid: str, debug: bool = False): +def show_pre_attach_info(server_pid: str, debug: bool = False) -> list: + """ + Collect pre-attach diagnostic information. + Returns a list of info messages to be printed only on failure. + """ from flight_profiler.utils.env_util import ( get_current_process_uids, get_process_uids, ) + messages = [] current_directory = os.path.dirname(os.path.abspath(__file__)) server_executable: str = get_py_bin_path(server_pid) client_executable: str = get_py_bin_path(os.getpid()) @@ -409,31 +791,31 @@ def show_pre_attach_info(server_pid: str, debug: bool = False): server_uids = get_process_uids(server_pid) client_uids = get_current_process_uids() - # Print executable information - print(f"PyFlightProfiler version: {version('flight_profiler')}") - print(f"[INFO] Platform system: {platform.system()}. Architecture: {platform.machine()}") - print(f"[INFO] Installation directory: {current_directory}.") + # Collect diagnostic information + messages.append(f"PyFlightProfiler version: {version('flight_profiler')}") + messages.append(f"[INFO] Platform system: {platform.system()}. Architecture: {platform.machine()}") + messages.append(f"[INFO] Installation directory: {current_directory}.") if debug: - print(f"[DEBUG] Server Python Executable: {server_executable}") - print(f"[DEBUG] Client Python Executable: {client_executable}") - print(f"[INFO] Verify pyFlightProfiler and target are using the same python executable: {'🌟' if same else '❌'}") + messages.append(f"[DEBUG] Server Python Executable: {server_executable}") + messages.append(f"[DEBUG] Client Python Executable: {client_executable}") + messages.append(f"[INFO] Verify pyFlightProfiler and target are using the same python executable: {'🌟' if same else '❌'}") # Check directory write permissions directory_write_permission = check_directory_write_permission(current_directory) permission_status = "🌟" if directory_write_permission else "❌" - print(f"[INFO] Verify pyFlightProfiler has write permission to installation directory: {permission_status}") + messages.append(f"[INFO] Verify pyFlightProfiler has write permission to installation directory: {permission_status}") if not directory_write_permission: - print(f"[WARN] PyFlightProfiler needs write permission to {current_directory} to function properly. " + messages.append(f"[WARN] PyFlightProfiler needs write permission to {current_directory} to function properly. " f"Please try run {COLOR_RED}flight_profiler with appropriate permissions{COLOR_END}.") - # Print permission information + # Collect permission information if server_uids and client_uids: server_real_uid, server_effective_uid, server_saved_uid, server_filesystem_uid = server_uids client_real_uid, client_effective_uid, client_saved_uid, client_filesystem_uid = client_uids if debug: - print(f"[INFO] Server Process - Real UID: {server_real_uid}, Effective UID: {server_effective_uid}") - print(f"[INFO] Client Process - Real UID: {client_real_uid}, Effective UID: {client_effective_uid}") + messages.append(f"[INFO] Server Process - Real UID: {server_real_uid}, Effective UID: {server_effective_uid}") + messages.append(f"[INFO] Client Process - Real UID: {client_real_uid}, Effective UID: {client_effective_uid}") # Check if client has sufficient privileges has_sufficient_privileges = ( @@ -443,15 +825,17 @@ def show_pre_attach_info(server_pid: str, debug: bool = False): ) privilege_status = "🌟" if has_sufficient_privileges else "❌" - print(f"[INFO] Verify pyFlightProfiler has user permission to attach target: {privilege_status}") + messages.append(f"[INFO] Verify pyFlightProfiler has user permission to attach target: {privilege_status}") # Additional check for root privileges if server_real_uid == 0 and client_effective_uid != 0: - print(f"[WARN] Target process is running as root, elevated privileges may be required.") + messages.append(f"[WARN] Target process is running as root, elevated privileges may be required.") elif client_effective_uid != 0 and server_real_uid != client_real_uid: - print(f"[WARN] Target process is owned by a different user, permission issues may occur.") + messages.append(f"[WARN] Target process is owned by a different user, permission issues may occur.") else: - print(f"[INFO] Permission information not available on this platform.") + messages.append(f"[INFO] Permission information not available on this platform.") + + return messages def run(): parser = argparse.ArgumentParser( @@ -473,7 +857,13 @@ def run(): inject_start_port = int(os.getenv("PYFLIGHT_INJECT_START_PORT", 16000)) inject_end_port = int(os.getenv("PYFLIGHT_INJECT_END_PORT", 16500)) inject_timeout = int(os.getenv("PYFLIGHT_INJECT_TIMEOUT", 5)) - show_pre_attach_info(server_pid, args.debug) + + # Collect diagnostic info (print immediately if debug mode, otherwise only on failure) + diagnostic_messages = show_pre_attach_info(server_pid, args.debug) + if args.debug: + for msg in diagnostic_messages: + print(msg) + print() # Empty line before welcome box connect_port: int = check_server_injected( server_pid, inject_start_port, inject_end_port, inject_timeout @@ -481,27 +871,34 @@ def run(): if connect_port < 0: free_port: int = find_port_available(inject_start_port, inject_end_port) if free_port < 0: + # Print diagnostic info on failure (skip if already printed in debug mode) + if not args.debug: + for msg in diagnostic_messages: + print(msg) print( f"No available debug port between range: {inject_start_port} {inject_end_port}" ) + exit(1) if sys.version_info >= (3, 14): if not is_linux() and not is_mac(): + if not args.debug: + for msg in diagnostic_messages: + print(msg) print(f"flight profiler is not enabled on platform: {platform.system()}.") exit(1) # sys.remote_exec is provided in CPython 3.14, we can just use it to inject agent code - connect_port = do_inject_with_sys_remote_exec(free_port, server_pid, args.debug) + connect_port = do_inject_with_sys_remote_exec(free_port, server_pid, args.debug, diagnostic_messages) else: if is_linux(): - connect_port = do_inject_on_linux(free_port, server_pid, args.debug) + connect_port = do_inject_on_linux(free_port, server_pid, args.debug, diagnostic_messages) elif is_mac(): - connect_port = do_inject_on_mac(free_port, server_pid, args.debug) + connect_port = do_inject_on_mac(free_port, server_pid, args.debug, diagnostic_messages) else: + if not args.debug: + for msg in diagnostic_messages: + print(msg) print(f"flight profiler is not enabled on platform: {platform.system()}.") exit(1) - else: - print( - f"[INFO] Process {server_pid} was attached through port {connect_port} already, so keep reusing the same port." - ) # add tab complete if READLINE_AVAILABLE: @@ -509,9 +906,11 @@ def run(): readline.parse_and_bind("tab: complete") cli = ProfilerCli(port=connect_port, target_executable=get_py_bin_path(server_pid)) check_preload = cli.check_status(timeout=5) - if check_preload: - print(f"\nPyFlightProfiler: 🌟 attach target process {server_pid} successfully!") - else: + if not check_preload: + # Print diagnostic info on failure (skip if already printed in debug mode) + if not args.debug: + for msg in diagnostic_messages: + print(msg) # here the injection routine is done successfully, but server has no chance to respond verify_exit_code(16, server_pid) diff --git a/flight_profiler/common/system_logger.py b/flight_profiler/common/system_logger.py index 0682bc5..d8efbeb 100644 --- a/flight_profiler/common/system_logger.py +++ b/flight_profiler/common/system_logger.py @@ -12,6 +12,12 @@ def setup_logger(level=None): Returns: logging.Logger: Configured logger instance """ + _logger = logging.getLogger("flight_profiler_logger") + + # Prevent duplicate handlers when module is imported multiple times + if _logger.handlers: + return _logger + if level is None: if os.getenv("FLIGHT_PROFILER_DEBUG", "0").strip().lower() in ("1", "true"): level = logging.DEBUG @@ -19,12 +25,13 @@ def setup_logger(level=None): level = logging.INFO handler = logging.StreamHandler() - formatter = logging.Formatter("%(asctime)s - %(filename)s - %(lineno)d - %(levelname)s - %(message)s") + formatter = logging.Formatter("%(asctime)s [PyFlightProfiler] %(levelname)s %(message)s") handler.setFormatter(formatter) - _logger = logging.getLogger("flight_profiler_logger") _logger.setLevel(level) _logger.addHandler(handler) + # Prevent propagation to root logger to avoid duplicate output + _logger.propagate = False return _logger diff --git a/flight_profiler/communication/flight_server.py b/flight_profiler/communication/flight_server.py index 7719ede..8ece4f1 100644 --- a/flight_profiler/communication/flight_server.py +++ b/flight_profiler/communication/flight_server.py @@ -71,8 +71,8 @@ async def handle_client(self, client_socket, addr): is_plugin_calling = request_json.get("is_plugin_calling", True) param = request_json.get("param", "") - logger.info( - f"[PyFlightProfiler] Cmd: {target} Param: {param} is_plugin_calling: {is_plugin_calling}" + logger.debug( + f"Cmd: {target} Param: {param} is_plugin_calling: {is_plugin_calling}" ) if is_plugin_calling: if target in self.interactive_commands: @@ -84,7 +84,7 @@ async def handle_client(self, client_socket, addr): else: await self.special_calling(target, param, writer) except: - logger.exception(f"[FlightServer] error in execute plugin") + logger.exception("error in execute plugin") finally: if writer is not None and not writer.is_closing(): writer.close() diff --git a/flight_profiler/help_descriptions.py b/flight_profiler/help_descriptions.py index e103add..55de14b 100644 --- a/flight_profiler/help_descriptions.py +++ b/flight_profiler/help_descriptions.py @@ -20,6 +20,7 @@ COLOR_WHITE_255, ICON_ARROW, align_prefix, + make_clickable_link, ) @@ -70,9 +71,10 @@ def _build_help_msg(self) -> str: wiki_description = "" if self._wiki is not None: + clickable_wiki = make_clickable_link(self._wiki) wiki_description = ( f"{section_prefix} WIKI{section_suffix}\n" - f" {COLOR_WHITE_255}{self._wiki}{COLOR_END}\n\n" + f" {COLOR_WHITE_255}{clickable_wiki}{COLOR_END}\n\n" ) return ( f"{section_prefix} USAGE{section_suffix}\n" diff --git a/flight_profiler/plugins/help/help_agent.py b/flight_profiler/plugins/help/help_agent.py index 336c04c..de7afb1 100644 --- a/flight_profiler/plugins/help/help_agent.py +++ b/flight_profiler/plugins/help/help_agent.py @@ -105,7 +105,7 @@ def display_all_commands(self): # Use ljust_display to handle emoji width correctly icon_part = ljust_display(cmd_icon, 2) name_part = ljust_display(cmd_name, 12) - display_msg += f"{COLOR_FAINT}{icon_part}{COLOR_END} {COLOR_BRIGHT_GREEN}{name_part}{COLOR_END}{align_prefix(15, HELP_COMMANDS_DESCRIPTIONS[c_idx].summary)}\n" + display_msg += f"{icon_part} {COLOR_BRIGHT_GREEN}{name_part}{COLOR_END}{align_prefix(15, HELP_COMMANDS_DESCRIPTIONS[c_idx].summary)}\n" return display_msg def get_command_description(self, command_name: str) -> str: diff --git a/flight_profiler/code_inject.py b/flight_profiler/profiler_agent.py similarity index 70% rename from flight_profiler/code_inject.py rename to flight_profiler/profiler_agent.py index 06d3aba..99d3063 100644 --- a/flight_profiler/code_inject.py +++ b/flight_profiler/profiler_agent.py @@ -24,18 +24,18 @@ def load_frida_gum(): nm_symbol_offset = int("${nm_symbol_offset}") flight_profiler_agent_so_path = "${flight_profiler_agent_so_path}" lib = ctypes.CDLL(flight_profiler_agent_so_path) - lib.inject_init_frida_gum.argtypes = [ctypes.c_ulong] - lib.inject_init_frida_gum.restype = ctypes.c_int - if lib.inject_init_frida_gum(nm_symbol_offset) != 0: - logger.warning(f"[PyFlightProfiler] init frid-gum failed, gilstat is disabled!") + lib.init_native_profiler.argtypes = [ctypes.c_ulong] + lib.init_native_profiler.restype = ctypes.c_int + if lib.init_native_profiler(nm_symbol_offset) != 0: + logger.warning("Native profiler init failed, gilstat is disabled!") except: - logger.exception(f"[PyFlightProfiler] flight_profiler_agent load failed!!!") + logger.exception("Native profiler agent load failed!") if PYTHON_VERSION_314: load_frida_gum() listen_port = int(listen_port) -logger.info("pyFlightProfiler: will use listen port " + str(listen_port)) +logger.debug(f"Agent listening on port {listen_port}") def run_app(): @@ -46,6 +46,6 @@ def run_app(): loop.run_until_complete(asyncio.wait(tasks)) -profile_thread = threading.Thread(target=run_app, name="flight-profiler-injector") +profile_thread = threading.Thread(target=run_app, name="flight-profiler-agent") profile_thread.start() -logger.info("pyFlightProfiler: start code inject successfully") +logger.debug("Agent thread started") diff --git a/flight_profiler/shell/code_inject.sh b/flight_profiler/shell/profiler_attach.sh similarity index 88% rename from flight_profiler/shell/code_inject.sh rename to flight_profiler/shell/profiler_attach.sh index 8acc6e2..b6dec63 100755 --- a/flight_profiler/shell/code_inject.sh +++ b/flight_profiler/shell/profiler_attach.sh @@ -40,28 +40,33 @@ debug_print() { fi } +# Error print function that always prints (for critical errors) +error_print() { + echo "$1" +} + py_bin_path=$(sh $shell_bin_dir/resolve_bin_path.sh $pid) if [ -z "$py_bin_path" ]; then - debug_print "Target process_id $pid not exists!" + error_print "Target process_id $pid not exists!" exit 1 fi client_py_bin_path=$(sh $shell_bin_dir/resolve_bin_path.sh $client_py_pid) if [ -z "$client_py_bin_path" ]; then - debug_print "client_py_pid $client_py_pid not exists" + error_print "client_py_pid $client_py_pid not exists" exit 1 fi if [ "$py_bin_path" != "$client_py_bin_path" ]; then - debug_print "target process does not use same python, profile client use $client_py_bin_path but target process use $py_bin_path" + error_print "target process does not use same python, profile client use $client_py_bin_path but target process use $py_bin_path" exit 1 fi nm_addr_hex=$(sh $shell_bin_dir/resolve_symbol.sh $pid take_gil) if [ -z "$nm_addr_hex" ]; then - debug_print "invalid python process $pid, test find take_gil function failed" + error_print "invalid python process $pid, test find take_gil function failed" exit 1 fi nm_addr=$(printf "%d" 0x${nm_addr_hex}) @@ -75,7 +80,7 @@ dylib="$parentdir/lib/flight_profiler_agent.${SHARED_LIB_SUFFIX}" echo "flight_profiler_agent.${SHARED_LIB_SUFFIX} not found." >&2 && \ echo "compile the library first." >&2 && exit 1 -pycode="$parentdir/code_inject.py" +pycode="$parentdir/profiler_agent.py" if [ "$IS_DARWIN" = "false" ]; then tmp_file=$(mktemp -p /tmp "$(basename $0).XXXXXX") @@ -101,7 +106,7 @@ if [ "$IS_DARWIN" = "false" ]; then echo "$gdb_logging_set_prefix off" echo "set \$m = (void*)dlopen(\"$dylib\", 9)" echo "set \$take_gil_addr = (void *)take_gil" - echo "set \$f = (int (*)(char *, int, unsigned long))dlsym(\$m, \"inject\")" + echo "set \$f = (int (*)(char *, int, unsigned long))dlsym(\$m, \"profiler_attach\")" echo "set \$used_port = \$f(\"$pycode\", $port_found, (unsigned long)\$take_gil_addr - $nm_addr)" echo "set logging overwrite on" echo "set logging file $tmp_file" @@ -116,7 +121,7 @@ else echo "process attach -p ${pid}" echo "expr void (* \$take_gil_addr)(void*) =(void (*)(void*))take_gil" echo "expr void* \$handle = (void*)dlopen(\"$dylib\", 9)" - echo "expr int (* \$func)(char*, int, unsigned long) = (int (*)(char*, int, unsigned long))dlsym(\$handle, \"inject\")" + echo "expr int (* \$func)(char*, int, unsigned long) = (int (*)(char*, int, unsigned long))dlsym(\$handle, \"profiler_attach\")" echo "expr int \$used_port = \$func(\"$pycode\", $port_found, (unsigned long)\$take_gil_addr - $nm_addr)" echo "print \$used_port" echo "c" diff --git a/flight_profiler/utils/render_util.py b/flight_profiler/utils/render_util.py index 610ad62..f3727fa 100644 --- a/flight_profiler/utils/render_util.py +++ b/flight_profiler/utils/render_util.py @@ -39,6 +39,47 @@ BANNER_COLOR_PURPLE ] + +def make_clickable_link(url: str, text: str = None) -> str: + """ + Create a clickable hyperlink for terminal using OSC 8 escape sequence. + Works in modern terminals like iTerm2, GNOME Terminal, Windows Terminal, etc. + Falls back to plain URL for unsupported terminals (e.g., Mac Terminal.app). + + Args: + url: The URL to link to + text: Display text (defaults to URL if not provided) + + Returns: + Terminal escape sequence for clickable link, or plain URL if unsupported + """ + if text is None: + text = url + + # Check if terminal supports OSC 8 hyperlinks + term_program = os.environ.get('TERM_PROGRAM', '') + term = os.environ.get('TERM', '') + + # Known terminals that support OSC 8 + supported_terminals = ['iTerm.app', 'vscode', 'WezTerm', 'Hyper'] + supported_terms = ['xterm-256color', 'screen-256color'] + + # Check for iTerm2, VSCode, or other known supporting terminals + is_supported = ( + term_program in supported_terminals or + 'ITERM_SESSION_ID' in os.environ or # iTerm2 + 'WT_SESSION' in os.environ or # Windows Terminal + 'KONSOLE_VERSION' in os.environ or # Konsole + 'GNOME_TERMINAL_SCREEN' in os.environ # GNOME Terminal + ) + + if is_supported: + # OSC 8 format: \033]8;;URL\007TEXT\033]8;;\007 + return f"\033]8;;{url}\007{text}\033]8;;\007" + else: + # Plain URL for unsupported terminals (most terminals auto-detect URLs) + return text + import unicodedata """ Status Icons @@ -86,6 +127,11 @@ BOX_CROSS = "┼" BOX_DOUBLE_HORIZONTAL = "═" BOX_LIGHT_HORIZONTAL = "╌" +# Rounded corners +BOX_ROUND_TOP_LEFT = "╭" +BOX_ROUND_TOP_RIGHT = "╮" +BOX_ROUND_BOTTOM_LEFT = "╰" +BOX_ROUND_BOTTOM_RIGHT = "╯" ENTRANCE_HINTS = [ @@ -337,6 +383,110 @@ def build_title_hints(additional_hints: List[Tuple[str, str]] = None) -> None: print() +def build_welcome_box(pid: str, py_executable: str) -> None: + """ + Build and display a Claude Code style welcome box. + + Args: + pid: Target process ID + py_executable: Path to Python executable + """ + ver = version("flight_profiler") + terminal_width = shutil.get_terminal_size().columns + box_width = min(80, terminal_width - 2) + + # Read banner + file_path = os.path.abspath(__file__) + dir_path = os.path.dirname(os.path.dirname(file_path)) + with open(os.path.join(dir_path, "banner.desc"), "r") as f: + banner_lines = f.read().splitlines() + + # Border color - use faint/gray for a subtle look + border_color = COLOR_FAINT + + # Build the box - title with white highlight and gray version number + title = "PyFlightProfiler" + version_str = f"v{ver}" + # Title: white bold (need COLOR_END first to clear FAINT attribute), version: gray + title_part = f"{COLOR_END}{COLOR_WHITE_255}{COLOR_BOLD}{title}{COLOR_END}" + version_part = f"{border_color}{version_str}" + title_display_len = len(f" {title} {version_str} ") + left_padding = 3 + right_padding = box_width - 2 - left_padding - title_display_len + top_line = f"{border_color}{BOX_ROUND_TOP_LEFT}{BOX_HORIZONTAL * left_padding} {title_part} {version_part} {BOX_HORIZONTAL * right_padding}{BOX_ROUND_TOP_RIGHT}{COLOR_END}" + + print() # Empty line before welcome box + print(top_line) + + # Helper function to print a boxed line + def print_box_line(content: str, content_display_len: int): + padding = box_width - 2 - content_display_len + print(f"{border_color}{BOX_VERTICAL}{COLOR_END}{content}{' ' * max(0, padding)}{border_color}{BOX_VERTICAL}{COLOR_END}") + + # Empty line + print_box_line("", 0) + + # Banner lines with colors + space_indices = [0] + if banner_lines: + for idx in range(1, len(banner_lines[0])): + all_space = True + for j in range(len(banner_lines)): + if idx < len(banner_lines[j]) and banner_lines[j][idx] != " ": + all_space = False + break + if all_space: + space_indices.append(idx) + + for line in banner_lines: + rendered_line = " " + for idx in range(len(space_indices)): + if space_indices[idx] >= len(line): + continue + if idx < len(space_indices) - 1: + rendered_line += f"{BANNER_COLOR_LIST[idx]}{COLOR_BOLD}{line[space_indices[idx]:space_indices[idx + 1]]}{COLOR_END}" + else: + rendered_line += f"{BANNER_COLOR_LIST[idx]}{COLOR_BOLD}{line[space_indices[idx]:]}{COLOR_END}" + display_len = str_display_width(line) + 2 + print_box_line(rendered_line, display_len) + + # Empty line + print_box_line("", 0) + + # Info section - each item on separate line + wiki_url = "https://github.com/alibaba/PyFlightProfiler/wiki" + clickable_wiki = make_clickable_link(wiki_url) + info_items = [ + ("pid", pid), + ("wiki", clickable_wiki), + ("python", py_executable), + ] + + for key, value in info_items: + line = f" {COLOR_FAINT}{key}:{COLOR_END} {COLOR_WHITE_255}{value}{COLOR_END}" + display_len = len(f" {key}: {value}") + print_box_line(line, display_len) + + # Empty line + print_box_line("", 0) + + # Bottom line + bottom_line = f"{border_color}{BOX_ROUND_BOTTOM_LEFT}{BOX_HORIZONTAL * (box_width - 2)}{BOX_ROUND_BOTTOM_RIGHT}{COLOR_END}" + print(bottom_line) + print() + + +def build_prompt_separator() -> str: + """ + Build a full-width separator line for the command prompt. + + Returns: + str: A separator line that spans terminal width + """ + terminal_width = shutil.get_terminal_size().columns + return f"{COLOR_FAINT}{BOX_HORIZONTAL * terminal_width}{COLOR_END}" + + def render_expression_result(result: ExpressionResult) -> str: left_offset: int = 20 diff --git a/pyproject.toml b/pyproject.toml index bca6739..df64777 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "flight_profiler" -version = "1.0.4" +version = "1.0.5" description = "A diagnostic toolkit for on-the-fly analysis of remote python processes." homepage = "https://github.com/alibaba/PyFlightProfiler" readme = "README.md"