diff --git a/README.md b/README.md index c00694e..8d5e98c 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ Generate eBPF programs and tracing with ChatGPT and natural language - Generate eBPF programs with natural language - +![generate](doc/generate.png) ## Usage diff --git a/doc/generate.png b/doc/generate.png new file mode 100644 index 0000000..72c7d9b Binary files /dev/null and b/doc/generate.png differ diff --git a/main.py b/main.py index 1a940b6..e847a94 100755 --- a/main.py +++ b/main.py @@ -25,17 +25,13 @@ def main(): group.add_argument( "-e", "--execute", help="Generate commands using your input with ChatGPT, and run it", action="store", metavar="TEXT") group.add_argument( - "-g", "--execute", help="Generate eBPF programs using your input with ChatGPT", action="store", metavar="TEXT") + "-g", "--generate", help="Generate eBPF programs using your input with ChatGPT", action="store", metavar="TEXT") parser.add_argument( "-u", "--uuid", help=f"Conversion UUID to use, or passed through environment variable `{ENV_UUID}`") parser.add_argument("-t", "--access-token", help=f"ChatGPT access token, see `https://chat.openai.com/api/auth/session` or passed through `{ENV_ACCESS_TOKEN}`") args = parser.parse_args() - if (args.info or args.execute) is None: - parser.print_help() - return - access_token = args.access_token or os.environ.get(ENV_ACCESS_TOKEN, None) conv_uuid = args.uuid or os.environ.get(ENV_UUID, None) if access_token is None: @@ -55,7 +51,25 @@ def main(): # print(f"Command to run: {parsed}") print("Press Ctrl+C to stop the program....") os.system("sudo " + parsed) + elif args.generate is not None: + desc: str = args.generate + print("Sending query to ChatGPT: " + desc) + ret_val = generate_result( + chatbot, construct_generate_prompt(desc), conv_uuid, True) + # print(ret_val) + parsed = extract_code_blocks(ret_val) + # print(f"Command to run: {parsed}") + with open("generated.bpf.c", "w") as f: + for code in parsed: + f.write(code) + else: + parser.print_help() +def construct_generate_prompt(text: str) -> str: + return f'''You are now a translater from human language to {os.uname()[0]} eBPF programs. +Please write eBPF programs for me. +No explanation required, no instruction required, don't tell me how to compile and run. +What I want is a eBPF program for: {text}.''' def construct_running_prompt(text: str) -> str: return f'''You are now a translater from human language to {os.uname()[0]} shell bpftrace command. diff --git a/prompts/statsnoop.bt b/prompts/statsnoop.bt deleted file mode 100755 index 14e1c0f..0000000 --- a/prompts/statsnoop.bt +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env bpftrace -/* - * statsnoop Trace stat() syscalls. - * For Linux, uses bpftrace and eBPF. - * - * This traces the tracepoints for statfs(), statx(), newstat(), and - * newlstat(). These aren't the only the stat syscalls: if you are missing - * activity, you may need to add more variants. - * - * Also a basic example of bpftrace. - * - * USAGE: statsnoop.bt - * - * This is a bpftrace version of the bcc tool of the same name. - * - * Copyright 2018 Netflix, Inc. - * Licensed under the Apache License, Version 2.0 (the "License") - * - * 08-Sep-2018 Brendan Gregg Created this. - */ - -BEGIN -{ - printf("Tracing stat syscalls... Hit Ctrl-C to end.\n"); - printf("%-6s %-16s %3s %s\n", "PID", "COMM", "ERR", "PATH"); -} - -tracepoint:syscalls:sys_enter_statfs -{ - @filename[tid] = args->pathname; -} - -tracepoint:syscalls:sys_enter_statx, -tracepoint:syscalls:sys_enter_newstat, -tracepoint:syscalls:sys_enter_newlstat -{ - @filename[tid] = args->filename; -} - -tracepoint:syscalls:sys_exit_statfs, -tracepoint:syscalls:sys_exit_statx, -tracepoint:syscalls:sys_exit_newstat, -tracepoint:syscalls:sys_exit_newlstat -/@filename[tid]/ -{ - $ret = args->ret; - $errno = $ret >= 0 ? 0 : - $ret; - - printf("%-6d %-16s %3d %s\n", pid, comm, $errno, - str(@filename[tid])); - delete(@filename[tid]); -} - -END -{ - clear(@filename); -}