|
|
@@ -1,43 +1,76 @@ |
|
|
#! /bin/env python3 |
|
|
#! /bin/env python3 |
|
|
import os |
|
|
|
|
|
import argparse |
|
|
import argparse |
|
|
|
|
|
|
|
|
from revChatGPT.V1 import Chatbot |
|
|
|
|
|
|
|
|
import os |
|
|
from typing import List |
|
|
from typing import List |
|
|
from marko.parser import Parser |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import pygments |
|
|
from marko.block import FencedCode |
|
|
from marko.block import FencedCode |
|
|
from marko.inline import RawText |
|
|
from marko.inline import RawText |
|
|
|
|
|
from marko.parser import Parser |
|
|
|
|
|
from prompt_toolkit import print_formatted_text |
|
|
|
|
|
from prompt_toolkit.formatted_text import PygmentsTokens |
|
|
|
|
|
from pygments.token import Token |
|
|
|
|
|
from pygments_markdown_lexer import MarkdownLexer |
|
|
|
|
|
from revChatGPT.V1 import Chatbot |
|
|
|
|
|
|
|
|
ENV_UUID = "GPTTRACE_CONV_UUID" |
|
|
ENV_UUID = "GPTTRACE_CONV_UUID" |
|
|
ENV_ACCESS_TOKEN = "GPTTRACE_ACCESS_TOKEN" |
|
|
ENV_ACCESS_TOKEN = "GPTTRACE_ACCESS_TOKEN" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def pretty_print(input, lexer=MarkdownLexer): |
|
|
|
|
|
tokens = list(pygments.lex(input, lexer=lexer())) |
|
|
|
|
|
print_formatted_text(PygmentsTokens(tokens)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main(): |
|
|
def main(): |
|
|
parser = argparse.ArgumentParser( |
|
|
parser = argparse.ArgumentParser( |
|
|
prog='GPTtrace', |
|
|
|
|
|
description='Use ChatGPT to write eBPF programs (bpftrace, etc.)') |
|
|
|
|
|
|
|
|
prog="GPTtrace", |
|
|
|
|
|
description="Use ChatGPT to write eBPF programs (bpftrace, etc.)", |
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
group = parser.add_mutually_exclusive_group() |
|
|
group = parser.add_mutually_exclusive_group() |
|
|
group.add_argument( |
|
|
group.add_argument( |
|
|
"-i", "--info", help="Let ChatGPT explain what's eBPF", action="store_true") |
|
|
|
|
|
|
|
|
"-i", "--info", help="Let ChatGPT explain what's eBPF", action="store_true" |
|
|
|
|
|
) |
|
|
group.add_argument( |
|
|
group.add_argument( |
|
|
"-e", "--execute", help="Generate commands using your input with ChatGPT, and run it", action="store", metavar="TEXT") |
|
|
|
|
|
|
|
|
"-e", |
|
|
|
|
|
"--execute", |
|
|
|
|
|
help="Generate commands using your input with ChatGPT, and run it", |
|
|
|
|
|
action="store", |
|
|
|
|
|
metavar="TEXT", |
|
|
|
|
|
) |
|
|
group.add_argument( |
|
|
group.add_argument( |
|
|
"-g", "--generate", 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( |
|
|
|
|
|
"-v", |
|
|
|
|
|
"--verbose", |
|
|
|
|
|
help="Print the prompt and receive message", |
|
|
|
|
|
action="store_true", |
|
|
|
|
|
) |
|
|
parser.add_argument( |
|
|
parser.add_argument( |
|
|
"-v", "--verbose", help="Print the prompt and receive message", action="store_true") |
|
|
|
|
|
|
|
|
"-u", |
|
|
|
|
|
"--uuid", |
|
|
|
|
|
help=f"Conversion UUID to use, or passed through environment variable `{ENV_UUID}`", |
|
|
|
|
|
) |
|
|
parser.add_argument( |
|
|
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}`") |
|
|
|
|
|
|
|
|
"-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() |
|
|
args = parser.parse_args() |
|
|
|
|
|
|
|
|
access_token = args.access_token or os.environ.get(ENV_ACCESS_TOKEN, None) |
|
|
access_token = args.access_token or os.environ.get(ENV_ACCESS_TOKEN, None) |
|
|
conv_uuid = args.uuid or os.environ.get(ENV_UUID, None) |
|
|
conv_uuid = args.uuid or os.environ.get(ENV_UUID, None) |
|
|
if access_token is None: |
|
|
if access_token is None: |
|
|
print( |
|
|
print( |
|
|
f"Either provide your access token through `-t` or through environment variable {ENV_ACCESS_TOKEN}") |
|
|
|
|
|
|
|
|
f"Either provide your access token through `-t` or through environment variable {ENV_ACCESS_TOKEN}" |
|
|
|
|
|
) |
|
|
return |
|
|
return |
|
|
chatbot = Chatbot(config={"access_token": access_token}) |
|
|
chatbot = Chatbot(config={"access_token": access_token}) |
|
|
if args.info: |
|
|
if args.info: |
|
|
@@ -46,7 +79,9 @@ def main(): |
|
|
desc: str = args.execute |
|
|
desc: str = args.execute |
|
|
print("Sending query to ChatGPT: " + desc) |
|
|
print("Sending query to ChatGPT: " + desc) |
|
|
ret_val = generate_result( |
|
|
ret_val = generate_result( |
|
|
chatbot, construct_running_prompt(desc), conv_uuid, args.verbose) |
|
|
|
|
|
|
|
|
chatbot, construct_running_prompt(desc), conv_uuid, args.verbose |
|
|
|
|
|
) |
|
|
|
|
|
pretty_print(ret_val) |
|
|
# print(ret_val) |
|
|
# print(ret_val) |
|
|
parsed = make_executable_command(ret_val) |
|
|
parsed = make_executable_command(ret_val) |
|
|
# print(f"Command to run: {parsed}") |
|
|
# print(f"Command to run: {parsed}") |
|
|
@@ -56,7 +91,8 @@ def main(): |
|
|
desc: str = args.generate |
|
|
desc: str = args.generate |
|
|
print("Sending query to ChatGPT: " + desc) |
|
|
print("Sending query to ChatGPT: " + desc) |
|
|
ret_val = generate_result( |
|
|
ret_val = generate_result( |
|
|
chatbot, construct_generate_prompt(desc), conv_uuid, True) |
|
|
|
|
|
|
|
|
chatbot, construct_generate_prompt(desc), conv_uuid, True |
|
|
|
|
|
) |
|
|
# print(ret_val) |
|
|
# print(ret_val) |
|
|
parsed = extract_code_blocks(ret_val) |
|
|
parsed = extract_code_blocks(ret_val) |
|
|
# print(f"Command to run: {parsed}") |
|
|
# print(f"Command to run: {parsed}") |
|
|
@@ -66,44 +102,47 @@ def main(): |
|
|
else: |
|
|
else: |
|
|
parser.print_help() |
|
|
parser.print_help() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def construct_generate_prompt(text: str) -> str: |
|
|
def construct_generate_prompt(text: str) -> str: |
|
|
return f'''You are now a translater from human language to {os.uname()[0]} eBPF programs. |
|
|
|
|
|
|
|
|
return f"""You are now a translater from human language to {os.uname()[0]} eBPF programs. |
|
|
Please write eBPF programs for me. |
|
|
Please write eBPF programs for me. |
|
|
No explanation required, no instruction required, don't tell me how to compile and run. |
|
|
No explanation required, no instruction required, don't tell me how to compile and run. |
|
|
What I want is a eBPF program for: {text}.''' |
|
|
|
|
|
|
|
|
What I want is a eBPF program for: {text}.""" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def construct_running_prompt(text: str) -> str: |
|
|
def construct_running_prompt(text: str) -> str: |
|
|
return f'''You are now a translater from human language to {os.uname()[0]} shell bpftrace command. |
|
|
|
|
|
|
|
|
return f"""You are now a translater from human language to {os.uname()[0]} shell bpftrace command. |
|
|
No explanation required. |
|
|
No explanation required. |
|
|
respond with only the raw shell bpftrace command. |
|
|
|
|
|
It should be start with `bpftrace`. |
|
|
|
|
|
|
|
|
respond with only the raw shell bpftrace command. |
|
|
|
|
|
It should be start with `bpftrace`. |
|
|
Your response should be able to put into a shell and run directly. |
|
|
Your response should be able to put into a shell and run directly. |
|
|
Just output in one line, without any description, or any other text that cannot be run in shell. |
|
|
Just output in one line, without any description, or any other text that cannot be run in shell. |
|
|
What should I type to shell to trace using bpftrace for: {text}, in one line.''' |
|
|
|
|
|
|
|
|
What should I type to shell to trace using bpftrace for: {text}, in one line.""" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def make_executable_command(command: str) -> str: |
|
|
def make_executable_command(command: str) -> str: |
|
|
if command.startswith('\n'): |
|
|
|
|
|
|
|
|
if command.startswith("\n"): |
|
|
command = command[1:] |
|
|
command = command[1:] |
|
|
if command.endswith('\n'): |
|
|
|
|
|
|
|
|
if command.endswith("\n"): |
|
|
command = command[:-1] |
|
|
command = command[:-1] |
|
|
if command.startswith('`'): |
|
|
|
|
|
|
|
|
if command.startswith("`"): |
|
|
command = command[1:] |
|
|
command = command[1:] |
|
|
if command.endswith('`'): |
|
|
|
|
|
|
|
|
if command.endswith("`"): |
|
|
command = command[:-1] |
|
|
command = command[:-1] |
|
|
command = command.strip() |
|
|
command = command.strip() |
|
|
command = command.split('User: ')[0] |
|
|
|
|
|
|
|
|
command = command.split("User: ")[0] |
|
|
return command |
|
|
return command |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_result(bot: Chatbot, text: str, session: str = None, print_out: bool = False) -> str: |
|
|
|
|
|
|
|
|
def generate_result( |
|
|
|
|
|
bot: Chatbot, text: str, session: str = None, print_out: bool = False |
|
|
|
|
|
) -> str: |
|
|
from io import StringIO |
|
|
from io import StringIO |
|
|
|
|
|
|
|
|
prev_text = "" |
|
|
prev_text = "" |
|
|
buf = StringIO() |
|
|
buf = StringIO() |
|
|
for data in bot.ask( |
|
|
|
|
|
text, conversation_id=session |
|
|
|
|
|
): |
|
|
|
|
|
message = data["message"][len(prev_text):] |
|
|
|
|
|
|
|
|
for data in bot.ask(text, conversation_id=session): |
|
|
|
|
|
message = data["message"][len(prev_text) :] |
|
|
if print_out: |
|
|
if print_out: |
|
|
print(message, end="", flush=True) |
|
|
print(message, end="", flush=True) |
|
|
buf.write(message) |
|
|
buf.write(message) |
|
|
|