Merge branch 'aicoder' of https://github.com/carlospolop/PEASS-ng into aicoder
This commit is contained in:
carlospolop
parent
d60fed0f20
commit
1279434ba6
208
AICoder.py
208
AICoder.py
@ -1,208 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import string
|
||||
import random
|
||||
from typing import List
|
||||
import openai
|
||||
import json
|
||||
import subprocess
|
||||
import tiktoken
|
||||
import requests
|
||||
from github import Github
|
||||
|
||||
#########################
|
||||
#### OPENAI FUNCTIONS ###
|
||||
#########################
|
||||
|
||||
def reportTokens(prompt, model="gpt-4"):
|
||||
encoding = tiktoken.encoding_for_model(model)
|
||||
print("\033[37m" + str(len(encoding.encode(prompt))) + " tokens\033[0m" + " in prompt: " + "\033[92m" + prompt[:50] + "\033[0m" + ("..." if len(prompt) > 50 else ""))
|
||||
|
||||
def write_file(file_path: str, content: str):
|
||||
"""Write content to a file creating the needed directories first"""
|
||||
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
||||
|
||||
with open(file_path, "w") as file:
|
||||
file.write(content)
|
||||
|
||||
def delete_file(file_path: str):
|
||||
"""Delete a file if it exists"""
|
||||
|
||||
if os.path.isfile(file_path):
|
||||
os.remove(file_path)
|
||||
|
||||
openai_available_functions = {
|
||||
"write_file": write_file, "delete_file": delete_file
|
||||
}
|
||||
|
||||
openai_functions = [
|
||||
{
|
||||
"name": "write_file",
|
||||
"description": "Write a file giving the path and the content",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the file to write",
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "Content to write in the file",
|
||||
},
|
||||
},
|
||||
"required": ["file_path", "content"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "delete_file",
|
||||
"description": "Delete a file",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"file_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the file to write",
|
||||
}
|
||||
},
|
||||
"required": ["file_path"],
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
#########################
|
||||
#### GIT FUNCTIONS ######
|
||||
#########################
|
||||
|
||||
|
||||
def create_pull_request(branch_name, commit_message, github_token):
|
||||
github = Github(github_token)
|
||||
repo = github.get_repo(os.environ["GITHUB_REPOSITORY"])
|
||||
|
||||
# Create a new branch
|
||||
base_branch = repo.get_branch(repo.default_branch)
|
||||
repo.create_git_ref(ref=f"refs/heads/{branch_name}", sha=base_branch.commit.sha)
|
||||
|
||||
# Commit changes to the new branch
|
||||
subprocess.run(["git", "checkout", branch_name])
|
||||
subprocess.run(["git", "add", "."])
|
||||
subprocess.run(["git", "commit", "-m", commit_message])
|
||||
subprocess.run(["git", "push", "origin", branch_name])
|
||||
|
||||
# Create a pull request
|
||||
pr = repo.create_pull(
|
||||
title=commit_message,
|
||||
body="Generated by OpenAI Github Action",
|
||||
head=branch_name,
|
||||
base=repo.default_branch
|
||||
)
|
||||
|
||||
return pr.html_url
|
||||
|
||||
|
||||
#########################
|
||||
#### FILE PROCESSING ####
|
||||
#########################
|
||||
|
||||
|
||||
def process_file(prompt: str, api_key: str, file_path: str, model: str="gpt-4") -> str:
|
||||
with open(file_path, "r") as file:
|
||||
file_content = file.read()
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": f"You are a developer and your goal is to generate code. The user will ask you to improve and modify some code. Your response must be a valid JSON with the path of each file to write as keys and the content of the files as values. Several files can be written at the same time."},
|
||||
{"role": "user", "content": prompt},
|
||||
{"role": "user", "content": f"This is the code from the file '{file_path}':\n\n{file_content}"}
|
||||
]
|
||||
openai.api_key = api_key
|
||||
|
||||
reportTokens(f"This is the code from the file '{file_path}':\n\n{file_content}")
|
||||
|
||||
response = openai.ChatCompletion.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0
|
||||
)
|
||||
response_message = response["choices"][0]["message"]
|
||||
|
||||
# Step 2: check if GPT wanted to call a function
|
||||
if response_message.get("function_call"):
|
||||
|
||||
function_name = response_message["function_call"]["name"]
|
||||
fuction_to_call = openai_available_functions[function_name]
|
||||
function_args = json.loads(response_message["function_call"]["arguments"])
|
||||
fuction_to_call(**function_args)
|
||||
|
||||
|
||||
def process_folder(prompt: str, api_key: str, folder_path: str, model: str="gpt-4") -> List[str]:
|
||||
responses = []
|
||||
for root, _, files in os.walk(folder_path):
|
||||
for file in files:
|
||||
file_path = os.path.join(root, file)
|
||||
response = process_file(prompt, api_key, file_path, model)
|
||||
responses.append(response)
|
||||
|
||||
|
||||
#########################
|
||||
#### MAIN FUNCTION ######
|
||||
#########################
|
||||
|
||||
|
||||
def get_random_string(length):
|
||||
# With combination of lower and upper case
|
||||
letters = string.ascii_letters
|
||||
result_str = ''.join(random.choice(letters) for i in range(length))
|
||||
return result_str
|
||||
|
||||
def main(prompt: str, api_key: str, file_path: str, github_token: str, model: str="gpt-4"):
|
||||
if os.path.isfile(file_path):
|
||||
process_file(prompt, api_key, file_path, model)
|
||||
elif os.path.isdir(file_path):
|
||||
process_folder(prompt, api_key, file_path, model)
|
||||
else:
|
||||
print("Error: Invalid file path.")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
create_pull_request(get_random_string(5), f"Modified {file_path}", github_token)
|
||||
except Exception as e:
|
||||
print(f"Error: Failed to create pull request. {e}")
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Setup the argument parser
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
# Add arguments for prompt, api_key, file_path and github_token
|
||||
parser.add_argument('--prompt', default=None, type=str, help='Input prompt')
|
||||
parser.add_argument('--api-key', default=None, type=str, help='Input API key')
|
||||
parser.add_argument('--path', default=None, type=str, help='Input file/folder path')
|
||||
parser.add_argument('--github-token', default=None, type=str, help='Github token')
|
||||
parser.add_argument('--model', default="gpt-4", type=str, help='Model to use')
|
||||
|
||||
# Parse the arguments
|
||||
args = parser.parse_args()
|
||||
prompt = os.environ.get("INPUT_PROMPT", args.prompt)
|
||||
api_key = os.environ.get("INPUT_API_KEY", args.api_key)
|
||||
file_path = os.environ.get("INPUT_FILE_PATH", args.path)
|
||||
github_token = os.environ.get("GITHUB_TOKEN", args.github_token)
|
||||
model = os.environ.get("INPUT_MODEL", args.model)
|
||||
|
||||
if not prompt or not api_key or not file_path:
|
||||
print("Error: Missing required inputs.")
|
||||
sys.exit(1)
|
||||
|
||||
#if not github_token:
|
||||
# print("Error: Missing github token.")
|
||||
# sys.exit(1)
|
||||
|
||||
if os.path.exists(prompt):
|
||||
with open(prompt, "r") as file:
|
||||
prompt = file.read()
|
||||
|
||||
if prompt.startswith("http"):
|
||||
prompt = requests.get(prompt).text
|
||||
|
||||
main(prompt, api_key, file_path, github_token, model)
|
||||
File diff suppressed because one or more lines are too long
@ -454,7 +454,7 @@ else
|
||||
sh_usrs=$(cat /etc/passwd 2>/dev/null | grep -v "^root:" | grep -i "sh$" | cut -d ":" -f 1 | tr '\n' '|' | sed 's/|bin|/|bin[\\\s:]|^bin$|/' | sed 's/|sys|/|sys[\\\s:]|^sys$|/' | sed 's/|daemon|/|daemon[\\\s:]|^daemon$|/')"ImPoSSssSiBlEee" #Modified bin, sys and daemon so they are not colored everywhere
|
||||
nosh_usrs=$(cat /etc/passwd 2>/dev/null | grep -i -v "sh$" | sort | cut -d ":" -f 1 | tr '\n' '|' | sed 's/|bin|/|bin[\\\s:]|^bin$|/')"ImPoSSssSiBlEee"
|
||||
fi
|
||||
knw_usrs='_amavisd|_analyticsd|_appinstalld|_appleevents|_applepay|_appowner|_appserver|_appstore|_ard|_assetcache|_astris|_atsserver|_avbdeviced|_calendar|_captiveagent|_ces|_clamav|_cmiodalassistants|_coreaudiod|_coremediaiod|_coreml|_ctkd|_cvmsroot|_cvs|_cyrus|_datadetectors|_demod|_devdocs|_devicemgr|_diskimagesiod|_displaypolicyd|_distnote|_dovecot|_dovenull|_dpaudio|_driverkit|_eppc|_findmydevice|_fpsd|_ftp|_fud|_gamecontrollerd|_geod|_hidd|_iconservices|_installassistant|_installcoordinationd|_installer|_jabber|_kadmin_admin|_kadmin_changepw|_knowledgegraphd|_krb_anonymous|_krb_changepw|_krb_kadmin|_krb_kerberos|_krb_krbtgt|_krbfast|_krbtgt|_launchservicesd|_lda|_locationd|_logd|_lp|_mailman|_mbsetupuser|_mcxalr|_mdnsresponder|_mobileasset|_mysql|_nearbyd|_netbios|_netstatistics|_networkd|_nsurlsessiond|_nsurlstoraged|_oahd|_ondemand|_postfix|_postgres|_qtss|_reportmemoryexception|_rmd|_sandbox|_screensaver|_scsd|_securityagent|_softwareupdate|_spotlight|_sshd|_svn|_taskgated|_teamsserver|_timed|_timezone|_tokend|_trustd|_trustevaluationagent|_unknown|_update_sharing|_usbmuxd|_uucp|_warmd|_webauthserver|_windowserver|_www|_wwwproxy|_xserverdocs|daemon\W|^daemon$|message\+|syslog|www|www-data|mail|noboby|Debian\-\+|rtkit|systemd\+'
|
||||
knw_usrs='_amavisd|_analyticsd|_appinstalld|_appleevents|_applepay|_appowner|_appserver|_appstore|_ard|_assetcache|_astris|_atsserver|_avbdeviced|_calendar|_captiveagent|_ces|_clamav|_cmiodalassistants|_coreaudiod|_coremediaiod|_coreml|_ctkd|_cvmsroot|_cvs|_cyrus|_datadetectors|_demod|_devdocs|_devicemgr|_diskimagesiod|_displaypolicyd|_distnote|_dovecot|_dovenull|_dpaudio|_driverkit|_eppc|_findmydevice|_fpsd|_ftp|_fud|_gamecontrollerd|_geod|_hidd|_iconservices|_installassistant|_installcoordinationd|_installer|_jabber|_kadmin_admin|_kadmin_changepw|_knowledgegraphd|_krb_anonymous|_krb_changepw|_krb_kadmin|_krb_kerberos|_krb_krbtgt|_krbfast|_krbtgt|_launchservicesd|_lda|_locationd|_logd|_lp|_mailman|_mbsetupuser|_mcxalr|_mdnsresponder|_mobileasset|_mysql|_nearbyd|_netbios|_netstatistics|_networkd|_nsurlsessiond|_nsurlstoraged|_oahd|_ondemand|_postfix|_postgres|_qtss|_reportmemoryexception|_rmd|_sandbox|_screensaver|_scsd|_securityagent|_softwareupdate|_spotlight|_sshd|_svn|_taskgated|_teamsserver|_timed|_timezone|_tokend|_trustd|_trustevaluationagent|_unknown|_update_sharing|_usbmuxd|_uucp|_warmd|_webauthserver|_windowserver|_www|_wwwproxy|_xserverdocs|daemon\W|^daemon$|message\+|syslog|www|www-data|mail|nobody|Debian\-\+|rtkit|systemd\+'
|
||||
if ! [ "$USER" ]; then
|
||||
USER=$(whoami 2>/dev/null || echo -n "UserUnknown")
|
||||
fi
|
||||
@ -1141,7 +1141,7 @@ if [ "$SEARCH_IN_FOLDER" ] || echo $CHECKS | grep -q procs_crons_timers_srvcs_so
|
||||
#GENERATE THE STORAGES OF THE FOUND FILES
|
||||
peass{STORAGES_HERE}
|
||||
|
||||
##### POST SERACH VARIABLES #####
|
||||
##### POST SEARCH VARIABLES #####
|
||||
backup_folders_row="$(echo $PSTORAGE_BACKUPS | tr '\n' ' ')"
|
||||
printf ${YELLOW}"DONE\n"$NC
|
||||
echo ""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user