├── .github └── ISSUE_TEMPLATE │ ├── bug.md │ ├── feature-request.md │ ├── documentation.md │ ├── inefficiency.md │ ├── ui-ux-problem.md │ └── security-weakness.md ├── requirements.txt ├── lib ├── banner.txt ├── config.py ├── classes.py └── functions.py ├── axiom ├── config.yml ├── LICENSE └── README.md /.github/ISSUE_TEMPLATE/bug.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug 3 | about: A bug in AXIOM Framework 4 | title: "[BUG]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature Request 3 | about: Request a reasonable new feature 4 | title: "[FEATURE]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/documentation.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Documentation 3 | about: Missing and/or inaccurate documentation 4 | title: "[DOCUMENTATION]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/inefficiency.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Inefficiency 3 | about: Recommend a better algorithm, implementation, method, etc. 4 | title: "[INEFFICIENCY]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/ui-ux-problem.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: UI/UX Problem 3 | about: A significant problem that negatively impacts usability for humans 4 | title: "[UI/UX]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | certifi==2019.11.28 2 | chardet==3.0.4 3 | colorama==0.4.3 4 | idna==2.9 5 | pexpect==4.8.0 6 | prompt-toolkit==3.0.3 7 | ptyprocess==0.6.0 8 | PyYAML==5.3 9 | requests==2.23.0 10 | urllib3==1.25.8 11 | wcwidth==0.1.8 12 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/security-weakness.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Security Weakness 3 | about: A vulnerability or other security weakness not addressed in README.md 4 | title: "[SECURITY]" 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /lib/banner.txt: -------------------------------------------------------------------------------- 1 | 2 | █████╗ ██╗ ██╗ ██╗ ██████╗ ███╗ ███╗ 3 | ██╔══██╗ ╚██╗██╔╝ ██║ ██╔═══██╗ ████╗ ████║ 4 | ███████║ ╚███╔╝ ██║ ██║ ██║ ██╔████╔██║ 5 | ██╔══██║ ██╔██╗ ██║ ██║ ██║ ██║╚██╔╝██║ 6 | ██║ ██║ ██╔╝ ██╗ ██║ ╚██████╔╝ ██║ ╚═╝ ██║ 7 | ╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═╝ 8 | 9 | 10 | -------------------------------------------------------------------------------- /axiom: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # Copyright 2020 Mike Iacovacci 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | from lib.functions import * 18 | 19 | 20 | def main(): 21 | 22 | try: 23 | settings = get_args() 24 | 25 | validate_privileges(settings.get("mode")) 26 | 27 | if settings.get("mode") == "init": 28 | initialize(settings) 29 | 30 | setup_folders(settings) 31 | 32 | if settings.get("mode") == "reload": 33 | reload() 34 | 35 | inventory = load_inventory() 36 | tool_list = load_tool_list(inventory) 37 | tool_names = get_tool_names(tool_list) 38 | tools = load_tools(inventory, tool_list[:]) 39 | 40 | branch(settings, tool_list, tools) 41 | 42 | print_stats(inventory, tool_list, tools) 43 | print_banner(config.axiom.banner_file) 44 | 45 | exit_code = axiom_prompt(tool_list, tool_names, tools) 46 | 47 | print("Exiting...") 48 | exit(exit_code) 49 | 50 | except KeyboardInterrupt: 51 | print_error("Keyboard interrupt received") 52 | return 1 53 | 54 | 55 | if __name__ == '__main__': 56 | main() 57 | 58 | -------------------------------------------------------------------------------- /config.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Mike Iacovacci 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | --- 15 | ### Which folder contains internal binary object files? 16 | binary_folder: ".bin" 17 | 18 | ### Which folder contains command input history files? 19 | history_folder: ".history" 20 | 21 | ### Which folder contains toolkit sub-folders where YAML files are stored? 22 | inventory_folder: "inventory" 23 | 24 | ### From which folder will The PenTesters Framework (PTF) be executed? 25 | ptf_folder: ".ptf" 26 | 27 | ### How many seconds will AXIOM Framework wait for new output after detecting an interactive subprogram's prompt? 28 | pattern_timeout: 1 29 | 30 | ### How many seconds will AXIOM Framework wait for new output before sending "ENTER key" to an interactive subprogram? 31 | safety_timeout: 15 32 | 33 | ### How many seconds will pexpect pseudo-terminal subprocesses wait before throwing any TIMEOUT exceptions? 34 | pty_timeout: 0.001 35 | 36 | ### AXIOM Framework will attempt to download and extract the items listed in "toolkits" when 1) the "inventory_folder" 37 | ### directory is missing and 2) the user explicitly initializes AXIOM Framework. During extraction the listed toolkit 38 | ### names are used to name sub-folders created within the inventory folder. AXIOM Framework expects toolkit data in ZIP 39 | ### format to be available at an HTTP(S) URL. The "file" attribute's value must match the name of the root folder within 40 | ### the ZIP archive. AXIOM Framework does not search for YAML data in nested sub-folders within a toolkit's directory. 41 | toolkits: 42 | - "Demo Toolkit X": 43 | - file: "axiom-data-demo-x-master" 44 | - url: "https://github.com/mikeiacovacci/axiom-data-demo-x/archive/master.zip" 45 | - "Demo Toolkit Y": 46 | - file: "axiom-data-demo-y-master" 47 | - url: "https://github.com/mikeiacovacci/axiom-data-demo-y/archive/master.zip" 48 | - "Demo Toolkit Z": 49 | - file: "axiom-data-demo-z-master" 50 | - url: "https://github.com/mikeiacovacci/axiom-data-demo-z/archive/master.zip" 51 | 52 | ### AXIOM Framework utilizes regex pattern matching for detecting interactive subprogram prompts. Regular expressions 53 | ### are modified at runtime to only match at the end of a subprogram output's line. All listed expressions must utilize 54 | ### backslash escaping, as needed, to be properly loaded by Python. Furthermore, prompt type names are case-sensitive. 55 | prompt_types: 56 | - bash: "[$#] " 57 | - cmd: "^[a-z]:(.+?)>" 58 | - empire: "^[(]Empire(.+?) > " 59 | - meterpreter: "\x1b\\[4mmeterpreter\x1b\\[0m > " 60 | - mimikatz: "mimikatz \\# " 61 | - msfconsole: "\x1b\\[4mmsf(.+?)> " 62 | - powershell: "^PS [A-Z]:(.+?)> " 63 | - rpcclient: "^rpcclient [$]> " 64 | - setoolkit: "\x1b\\[4m\x1b\\[36mset\x1b\\[0m(.+?)" 65 | - smbclient: "^smb: \\\\> " 66 | - other: null 67 | 68 | ### Which text file contains the AXIOM Framework logo? 69 | banner_file: "lib/banner.txt" 70 | 71 | ### AXIOM Framework uses values listed in "input_types" to 1) interpret command text by identifying the number and type 72 | ### of user inputs a command requires 2) inform the end user about the command input requirements, and 3) creating and 73 | ### updating history files to auto-suggest command inputs. Input type names must be less than or equal to eight 74 | ### characters. End users can add more types, but AXIOM Framework does not attempt any input data validation at all. 75 | input_types: 76 | - DOMAIN 77 | - FILE 78 | - FULLPATH 79 | - HTTPSURL 80 | - HTTPURL 81 | - INT 82 | - INTMENU 83 | - INTRANGE 84 | - IPV4 85 | - IPV4CIDR 86 | - IPV4RNGE 87 | - IPV6 88 | - IPV6CIDR 89 | - IPV6RNGE 90 | - MAC 91 | - RLATVPTH 92 | - STR 93 | - STRMENU 94 | - WEBURL 95 | 96 | ### All interactive commands must list a PROMPT output specifying the name of the prompt type (matching an item listed 97 | ### in "prompt_types") to be properly detected. AXIOM Framework currently only handles STDOUT and PROMPT output types. 98 | output_types: 99 | - F_INPUT 100 | - F_PREFIX 101 | - F_STRING 102 | - PROMPT 103 | - STDERR 104 | - STDOUT 105 | ... 106 | -------------------------------------------------------------------------------- /lib/config.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Mike Iacovacci 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from colorama import Fore, Style 16 | from os import path 17 | from sys import platform, stderr 18 | from yaml import parser, safe_load_all, scanner 19 | 20 | 21 | class AxiomConfig: 22 | """ a collection of elements the program relies on for data validation and proper execution """ 23 | 24 | def __init__(self, supplied_filename): 25 | """ creates multi-purpose, global config object used throughout program execution 26 | INPUT: the filename of the YAML configuration file, typically "config.yml" 27 | OUTPUT: instantiates an AxiomConfig object """ 28 | 29 | self.config_file = str(supplied_filename) 30 | self.yaml_list = self.get_yaml(self.config_file) 31 | 32 | self.platform = self.get_platform() 33 | 34 | self.binary_folder = None 35 | self.history_folder = None 36 | self.inventory_folder = None 37 | self.ptf_folder = None 38 | self.get_folders() 39 | 40 | self.pattern_timeout = None 41 | self.pty_timeout = None 42 | self.safety_timeout = None 43 | self.get_timeouts() 44 | 45 | self.toolkits = self.get_toolkits() 46 | 47 | self.prompts = self.get_prompts() 48 | 49 | self.banner_file = self.get_banner() 50 | 51 | self.inputs_pattern = None 52 | self.input_types_list = None 53 | self.get_inputs() 54 | 55 | self.outputs = self.get_outputs() 56 | 57 | def get_banner(self): 58 | """ validates user-supplied banner filename, returns a filename (str) """ 59 | 60 | try: 61 | banner = str(self.yaml_list[0]["banner_file"]) 62 | 63 | if banner in ["", "None"]: 64 | print_error("ERROR: Invalid banner_file setting in configuration file") 65 | exit(1) 66 | 67 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 68 | print_error("ERROR: Invalid banner_file setting in configuration file") 69 | exit(1) 70 | 71 | else: 72 | return banner 73 | 74 | def get_folders(self): 75 | """ validates user-supplied folder names and sets global config values """ 76 | 77 | try: 78 | binary_folder = str(self.yaml_list[0]["binary_folder"]) 79 | history_folder = str(self.yaml_list[0]["history_folder"]) 80 | inventory_folder = str(self.yaml_list[0]["inventory_folder"]) 81 | ptf_folder = str(self.yaml_list[0]["ptf_folder"]) 82 | 83 | if "None" in [binary_folder, history_folder, inventory_folder, ptf_folder] or \ 84 | "" in [binary_folder, history_folder, inventory_folder, ptf_folder]: 85 | print_error("ERROR: Invalid folder setting(s) in configuration file") 86 | exit(1) 87 | 88 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 89 | print_error("ERROR: Invalid folder setting(s) in configuration file") 90 | exit(1) 91 | 92 | else: 93 | self.binary_folder = binary_folder 94 | self.history_folder = history_folder 95 | self.inventory_folder = inventory_folder 96 | self.ptf_folder = ptf_folder 97 | 98 | @staticmethod 99 | def get_platform(): 100 | """ queries the system for platform type, returns a platform name (str) """ 101 | 102 | if platform == "linux": 103 | return "Linux" 104 | elif platform == "darwin": 105 | return "macOS" 106 | else: 107 | return "UNKNOWN PLATFORM" 108 | 109 | def get_inputs(self): 110 | """ iterates over listed input types, sets multiple values for class variables """ 111 | 112 | input_types_list = [] 113 | inputs_pattern = "" 114 | 115 | try: 116 | input_count = self.yaml_list[0]["input_types"].__len__() 117 | for i in range(input_count): 118 | input_name = str(self.yaml_list[0]["input_types"][i]) 119 | 120 | if input_name == "None": 121 | print_error("ERROR: Invalid input type in configuration file") 122 | exit(1) 123 | 124 | input_types_list.append(input_name) 125 | inputs_pattern = inputs_pattern + "{" + input_name + "}|" 126 | inputs_pattern = inputs_pattern[:-1] 127 | 128 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 129 | print_error("ERROR: Configuration file error(s) near input_types section") 130 | exit(1) 131 | 132 | else: 133 | self.inputs_pattern = inputs_pattern 134 | self.input_types_list = input_types_list 135 | 136 | def get_outputs(self): 137 | """ iterates over listed output types, returns a list of two-item tuples """ 138 | 139 | output_types = [] 140 | 141 | try: 142 | output_count = self.yaml_list[0]["output_types"].__len__() 143 | for i in range(output_count): 144 | output_name = str(self.yaml_list[0]["output_types"][i]) 145 | 146 | if output_name == "None": 147 | print_error("ERROR: Invalid output type in configuration file") 148 | exit(1) 149 | 150 | output_types.append(output_name) 151 | 152 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 153 | print_error("ERROR: Configuration file error(s) near output_types section") 154 | exit(1) 155 | 156 | else: 157 | return output_types 158 | 159 | def get_prompts(self): 160 | """ iterates over listed prompt types, returns a list of two-item tuples """ 161 | 162 | prompt_types = [] 163 | 164 | try: 165 | prompt_count = self.yaml_list[0]["prompt_types"].__len__() 166 | for i in range(prompt_count): 167 | prompt_name = list(self.yaml_list[0]["prompt_types"][i].keys())[0] 168 | prompt_pattern = str(self.yaml_list[0]["prompt_types"][i][prompt_name]) 169 | 170 | if prompt_name != "other" and prompt_pattern == "None": 171 | print_error(str("ERROR: Invalid prompt pattern for \"" + prompt_name + "\" in configuration file")) 172 | exit(1) 173 | 174 | current_prompt = (prompt_name, prompt_pattern) 175 | prompt_types.append(current_prompt) 176 | 177 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 178 | print_error("ERROR: Configuration file error(s) near prompt_types section") 179 | exit(1) 180 | 181 | else: 182 | return prompt_types 183 | 184 | def get_timeouts(self): 185 | """ validates user-supplied timeout values and sets them in the global config """ 186 | 187 | try: 188 | pattern_timeout = int(self.yaml_list[0]["pattern_timeout"]) 189 | pty_timeout = float(self.yaml_list[0]["pty_timeout"]) 190 | safety_timeout = int(self.yaml_list[0]["safety_timeout"]) 191 | 192 | if pattern_timeout < 0 or \ 193 | pty_timeout < 0 or \ 194 | safety_timeout < 0: 195 | print_error("ERROR: Invalid timeout setting(s) in configuration file") 196 | exit(1) 197 | 198 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 199 | print_error("ERROR: Invalid timeout setting(s) in configuration file") 200 | exit(1) 201 | 202 | else: 203 | self.pattern_timeout = pattern_timeout 204 | self.pty_timeout = pty_timeout 205 | self.safety_timeout = safety_timeout 206 | 207 | def get_toolkits(self): 208 | """ iterates over listed toolkits in the YAML file, returns a list of three-item tuples """ 209 | 210 | toolkits = [] 211 | 212 | try: 213 | toolkit_count = self.yaml_list[0]["toolkits"].__len__() 214 | for i in range(toolkit_count): 215 | toolkit_name = str(list(self.yaml_list[0]["toolkits"][i].keys())[0]) 216 | toolkit_file = str(self.yaml_list[0]["toolkits"][i][toolkit_name][0]["file"]) 217 | toolkit_url = str(self.yaml_list[0]["toolkits"][i][toolkit_name][1]["url"]) 218 | 219 | if (toolkit_name == "" or toolkit_name is None) or \ 220 | (toolkit_file == "" or toolkit_file is None) or \ 221 | (toolkit_url == "" or toolkit_url is None): 222 | print_error("ERROR: Undefined toolkit in configuration file") 223 | exit(1) 224 | else: 225 | if self.yaml_list[0]["toolkits"][i][toolkit_name].__len__() == 2: 226 | current_toolkit = (toolkit_name, toolkit_file, toolkit_url) 227 | toolkits.append(current_toolkit) 228 | else: 229 | print_error("ERROR: Invalid toolkit in configuration file") 230 | exit(1) 231 | 232 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 233 | print_error("ERROR: Configuration file error(s) near toolkits section") 234 | exit(1) 235 | 236 | else: 237 | return toolkits 238 | 239 | @staticmethod 240 | def get_yaml(config_file): 241 | """ extracts YAML content from specified file, returns a list object """ 242 | 243 | try: 244 | if not path.exists(config_file): 245 | print_error("ERROR: Missing configuration file") 246 | exit(1) 247 | 248 | with open(config_file, 'r') as open_file: 249 | yaml_list = list(safe_load_all(open_file)) 250 | 251 | except IOError: 252 | print_error("ERROR: Failed to open configuration file") 253 | exit(1) 254 | except (parser.ParserError, scanner.ScannerError): 255 | print_error("ERROR: Invalid configuration file") 256 | exit(1) 257 | else: 258 | return yaml_list 259 | 260 | 261 | def print_error(message): 262 | """ SUMMARY: prints stylized error text to STDERR 263 | INPUT: error message (str) 264 | OUTPUT: no return value, prints to STDERR """ 265 | 266 | stderr.write(Fore.RED + message + Style.RESET_ALL + "\n") 267 | 268 | 269 | axiom = AxiomConfig("config.yml") 270 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |

2 | 3 |
4 | 5 | 6 | Release 8 | 9 | 10 | 11 | License 13 | 14 | 15 | 16 | Issues 18 | 19 | 20 | 21 | Espresso 23 | 24 |
25 | 26 |
27 | 28 |

29 | AXIOM is a configurable, interactive knowledge management framework for learning, using, and experimenting with 30 | arbitrary command line programs. 31 |

32 | 33 | --- 34 | AXIOM Framework lets you "bookmark" your commands so you can reference, modify, and execute them more easily. 35 | 36 | If you know what you're doing then get started with the following: 37 | 38 | ``` 39 | git clone git@github.com:mikeiacovacci/axiom-framework.git 40 | cd axiom-framework 41 | pip3 install -r requirements.txt 42 | sudo python3 ./axiom 43 | ``` 44 | 45 | Otherwise, please read the [security](#security) considerations before installing. 46 | 47 | ## Table of Contents 48 | 49 | - [Motivation](#motivation) 50 | - [What it Does](#what-it-does) 51 | - [Features](#features) 52 | - [Implementation](#implementation) 53 | - [Installation](#installation) 54 | - [Dependencies](#dependencies) 55 | - [Security](#security) 56 | - [Usage](#usage) 57 | - [Standard Usage](#standard-usage) 58 | - [Referencing](#referencing) 59 | - [Modifying](#modifying) 60 | - [Executing](#executing) 61 | - [Interactive Programs](#interactive-programs) 62 | - [Configuration](#configuration) 63 | - [Adding Commands](#adding-commands) 64 | - [Known Limitations](#known-limitations) 65 | - [Feedback & Support](#feedback--support) 66 | - [License](#license) 67 | 68 | ## Motivation 69 | 70 | Infosec professionals are expected to know how to use **hundreds** of command line software programs. New offensive and 71 | defensive tools get released *all the time*, and practitioners need to learn them through hands-on experience. One might 72 | do any of the following to get started: 73 | 74 | - Read manual pages 75 | - Search online 76 | - Elicit program help/usage text 77 | - Trial and error 78 | 79 | Some command line programs present learning obstacles like missing documentation and poor feature discoverability. 80 | Furthermore, if one doesn't interact with a given CLI program for a while it's natural to forget the specific details 81 | (command syntax, available features, expected outputs, environmental requirements, "lesson's learned" from last time, 82 | etc.) needed to use the tool effectively. When this happens, one can **repeat** the above steps with the addition of two 83 | more: 84 | 85 | - Revisit terminal history 86 | - Refer to any personal notes 87 | 88 | These approaches definitely help with recall, but rereading one's notes (unstructured text, handwriting, etc.) or 89 | searching one's terminal history (if it exists on the current system) then copying/pasting, modifying, or retyping 90 | command text is an *annoying* distraction from the task at hand and wastes one's time and limited capacity to focus. 91 | 92 | Even when recall isn't a problem infosec professionals still need to *maximize* their **learning** while *also* 93 | improving their **routine tool use** by *minimizing* time-wasting activities (repetitive typing, manually installing 94 | tools, etc.) at the command line. 95 | 96 | ## What it Does 97 | 98 | ### Features 99 | 100 | - Presents an interactive, keystroke-conserving, and discoverability-oriented interface 101 | - Uses a prompt-driven, wizard-like menu system for creating, modifying, and running commands 102 | - Auto-suggests type-specific command inputs based on history 103 | - Executes commands via local subprocesses or by transmitting text to interactive prompts 104 | - Supports any interactive CLI program that outputs a prompt with a detectable pattern 105 | - Supports single- and multi-line standalone and interactive commands 106 | - Supports single-line "autonomous" commands (i.e. shell commands that use pipes) 107 | - Generates single- and multi-line non-executable text (e.g. injectable payloads) 108 | - Installs missing programs automatically via [The PenTesters Framework (PTF)](https://github.com/trustedsec/ptf) 109 | - Integrates command data from multiple "toolkits" hosted anywhere 110 | - Supports user-defined prompts, prompt patterns, and input data types 111 | - Supports deploying a custom config at first execution via URL parameter 112 | - Runs on Linux (Debian, Ubuntu, and ArchLinux) with partial macOS support 113 | - Utilizes a human-readable data format 114 | 115 | ### Implementation 116 | 117 | AXIOM Framework relies on YAML files that contain data about CLI programs, their commands, and details about those 118 | commands. Generally, each program is represented as a `.yml` file that specifies the program's name, a description, the 119 | operating system on which the program runs, and a PTF module if applicable. 120 | 121 | These YAML files can contain an infinite number of listed command entries that, in turn, specify the command's name 122 | (i.e. a brief description or "nickname"), prompt and execution types, text with any placeholder values, input 123 | descriptions, outputs, and any notes for the end user. 124 | 125 | AXIOM Framework integrates all the data at runtime by 126 | 127 | 1. searching the `inventory folder` (specified in the configuration file) for sub-folders ("toolkits") containing YAML 128 | files 129 | 2. merging the data from an infinite number of said YAML files into a unified structure. 130 | 131 | The program generally interacts with user-provided data in a read-only capacity, but it's built to fetch ZIP-compressed 132 | toolkits from any HTTP(S) URL to encourage the user to keep data separated (e.g. in a version control system) for 133 | improved loss-prevention and collaboration. 134 | 135 | ## Installation 136 | 137 | Install AXIOM Framework by cloning this repository or downloading and extracting the ZIP file in 138 | [Releases](https://github.com/mikeiacovacci/axiom-framework/releases). 139 | 140 | ### Dependencies 141 | 142 | - POSIX platform (Debian, Ubuntu, or ArchLinux for full compatibility) 143 | - Python 3 (including modules listed in `requirements.txt`) 144 | - `bash` 145 | - `which` 146 | 147 | Know that `axiom` utilizes a shebang of `#!/usr/bin/env python3` which may not work for a given `python3` installation. 148 | The end user can override it by supplying the `axiom` file as an argument to the Python 3 interpreter. 149 | 150 | ### Security 151 | 152 | End users are strongly advised to run AXIOM Framework on disposable, *untrusted* infrastructure after downloading the 153 | ZIP file in [Releases](https://github.com/mikeiacovacci/axiom-framework/releases) and verifying the PGP signature. **Do 154 | not run the framework on critical systems**, and don't use it *at all* if you don't trust me or my code :) 155 | 156 | Furthermore, know that a "good signature" only indicates the ZIP file contents were signed by the corresponding private 157 | PGP key and this does not validate the authenticity of any dependencies like third-party libraries/modules or toolkit 158 | datasets. AXIOM Framework downloads, extracts, installs, and executes third-party software and other content. **Do not 159 | use the framework if you don't trust any of those third parties.** 160 | 161 | Users should regard toolkit data as executable content since AXIOM Framework uses it 162 | 163 | 1. to execute shell commands 164 | 2. to spawn local subprocesses 165 | 3. as input transmitted to other programs 166 | 167 | Additionally, even before the one chooses to execute anything, the data is deserialized into Python objects, and 168 | malformed YAML input could hypothetically abuse program logic or execute arbitrary code. **Do not use toolkits from 169 | untrusted sources.** 170 | 171 | AXIOM Framework requires root privileges for many modes of operation. Whoever controls the hosting infrastructure for 172 | the toolkit data or any optional, custom config file could hypothetically achieve code execution as root on your 173 | machine. 174 | 175 | Lastly, the framework does not distinguish secret from non-secret user input. Any passwords, keys, or other sensitive 176 | inputs will be displayed on the screen and stored on disk, in plaintext, in the `history folder` (specified in the 177 | configuration file) within one or more `.axiom` history files. 178 | 179 | ## Usage 180 | 181 | ### Standard Usage 182 | 183 | To interact with AXIOM Framework run `./axiom` and follow the prompts by 184 | 185 | 1. entering the name of a tool 186 | 2. selecting a command by number 187 | 3. confirming execution 188 | 4. entering any required inputs 189 | 190 | ![AXIOM Framework executing nmap](https://payl0ad.run/assets/images/post-8/axiom-framework-nmap.gif "Executing a simple command interactively") 191 | 192 | This interface is useful for executing [interactive programs](#interactive-programs) and when switching between multiple 193 | tools. To select a different tool enter `back` at the command selection prompt. Entering `exit` at either the command or 194 | tool selection prompt will terminate the program. 195 | 196 | Users interested in a specific, non-interactive tool can supply the tool name as a command line argument. Tool names 197 | are case sensitive. A tool name that contains spaces must be passed as a singular argument by enclosing the entire name 198 | in quotes or backslash-escaping the space characters. 199 | 200 | ### Referencing 201 | 202 | To view information about a tool enter `./axiom show [TOOL]` supplying the tool name. AXIOM Framework will display the 203 | tool's PTF module (if any), notes, and an alphabetized list of the available commands. 204 | 205 | ![AXIOM Framework showing sqlmap](https://payl0ad.run/assets/images/post-8/axiom-framework-sqlmap.gif "AXIOM Framework listing a tool's commands") 206 | 207 | To see more details about a specific command enter `./axiom show [TOOL] [NUM]` providing the tool name *and* the command 208 | number. This prints the command's name, execution and prompt types, notes, and text showing the placeholder values. 209 | 210 | ![AXIOM Framework showing hashcat command](https://payl0ad.run/assets/images/post-8/axiom-framework-hashcat.gif "AXIOM Framework showing command details") 211 | 212 | ### Modifying 213 | 214 | To enter input values (i.e. to replace a command's placeholders) and print executable, "finalized" command text to the 215 | screen (e.g. to copy/paste into a script or another prompt) run `./axiom build [TOOL] [NUM]` with the tool name and 216 | command number and follow the prompts. If a command does not require user-supplied values then the text will simply be 217 | displayed on the screen. 218 | 219 | ![AXIOM Framework generating PowerShell text](https://payl0ad.run/assets/images/post-8/axiom-framework-powershell.gif "Building command text interactively") 220 | 221 | ### Executing 222 | 223 | Users can execute non-interactive commands locally via AXIOM Framework by running `./axiom run [TOOL] [NUM]` supplying 224 | the tool name and command number as CLI arguments. 225 | 226 | ![AXIOM Framework executing lsof](https://payl0ad.run/assets/images/post-8/axiom-framework-lsof.gif "Executing a standalone action non-interactively") 227 | 228 | If the command requires input values, the user will be prompted to enter them. Otherwise, the command will simply 229 | execute. Additionally, non-executable commands will merely print command text to the screen. 230 | 231 | ![AXIOM Framework running Python](https://payl0ad.run/assets/images/post-8/axiom-framework-python.gif "Outputting command text") 232 | 233 | ### Interactive Programs 234 | 235 | AXIOM Framework supports executing interactive subprograms by: 236 | 237 | 1. transmitting command text to pseudo-terminal subprocesses 238 | 2. detecting subprogram input prompts in STDOUT via regex pattern matching 239 | 3. observing user-configurable timeouts to reduce false positives 240 | 241 | When the user runs an interactive command, AXIOM Framework transmits the command text as subprogram input and prints the 242 | subprogram's STDOUT to the screen until the expected prompt pattern is detected. 243 | 244 | ![AXIOM Framework interacting with msfconsole](https://payl0ad.run/assets/images/post-8/axiom-framework-msfconsole.gif "A simple interactive program") 245 | 246 | Executing interactive commands can result in prompt changes, so 247 | [AXIOM Framework maintains state](https://payl0ad.run/assets/images/post-8/axiom-framework-multiple-prompts.gif) 248 | to ensure command text is transmitted to the correct pseudo-terminal. It also prevents the user from creating runtime 249 | ambiguities by blocking commands that result in more than one subprocess having identical prompt types. 250 | 251 | When running any executable command, AXIOM Framework **always** attempts to transmit command text (i.e. instead of 252 | executing commands locally) when the current runtime includes a pseudo-terminal subprocess with a matching prompt type. 253 | 254 | ## Configuration 255 | 256 | AXIOM Framework expects a file named `config.yml` in the top-level folder that modifies the program's interaction with 257 | the filesystem, subprocesses, toolkit data, and the end user. 258 | 259 | Launching `axiom` for the first time (without any command line arguments) initializes the framework with settings from 260 | the default `config.yml` file after creating any missing folders and downloading/extracting any missing content. 261 | 262 | To manually initialize (to reinstall PTF, re-download toolkits, etc.) run `./axiom init`. Initializing can result in 263 | data loss, because the folders listed in `config.yml` will be deleted and replaced. Folder settings can hypothetically 264 | reference directories outside of the top-level folder, so verify the settings are correct before executing. 265 | 266 | To manually initialize with a *custom* configuration run `./axiom init [URL]` specifying an HTTP(S) URL hosting the new 267 | configuration file. AXIOM Framework will 268 | 269 | 1. ignore the existing `config.yml` file 270 | 2. download the new file (replacing the existing one) 271 | 3. initialize with the new settings 272 | 273 | ![Customizing AXIOM Framework](https://payl0ad.run/assets/images/post-8/axiom-framework-custom-config.gif "Configuring AXIOM Framework with custom settings") 274 | 275 | Again, initializing can result in data loss, because the folders listed in the *new* configuration file will be deleted 276 | and replaced. Only deploy config files created by trusted parties and hosted on trusted infrastructure. 277 | 278 | To integrate any changes to the *local* YAML data (e.g. while learning and experimenting) run `./axiom reload`. AXIOM 279 | Framework will reprocess all the YAML files which could take a few additional seconds. Know that modifications to local 280 | YAML data will be *lost* if initialization occurs, so ensure that *permanent* changes are saved at the data source. 281 | Additionally, a tool's commands are always listed in alphabetical order, and modifying a command's name can change its 282 | ID number and command list ordering without warning. 283 | 284 | Worth noting is that neither initializing nor reloading causes AXIOM Framework to delete its `.axiom` history files. The 285 | end user can manually delete specific history files or the whole `history folder` to reset input auto-suggestion. 286 | 287 | ## Adding Commands 288 | 289 | AXIOM Framework is not merely the sum of any existing toolkit datasets. It's designed to be built upon by end users. 290 | Certainly toolkits can be shared, but users intending to create and use *their own custom toolkits* can write new YAML 291 | data "from scratch" with the help of two approaches: 292 | 293 | 1. Learn and borrow from more than **1,000** [examples](https://github.com/mikeiacovacci/axiom-data-demo-x) 294 | [available](https://github.com/mikeiacovacci/axiom-data-demo-y) 295 | [online](https://github.com/mikeiacovacci/axiom-data-demo-z). 296 | 2. Run `./axiom new` for an interactive "wizard" mode that generates valid YAML. 297 | 298 | ![AXIOM Framework generating YAML](https://payl0ad.run/assets/images/post-8/axiom-framework-new-command.gif "Generating new toolkit data interactively") 299 | 300 | Be aware of the following ideas, constraints, and best practices when creating custom toolkit data: 301 | 302 | - A tool can be represented via multiple, separate YAML files across more than one toolkit. 303 | - A tool with more than one YAML file must have matching `name`, `description`, `os`, and `ptf_module` values to merge. 304 | - A tool's `os` value affects if AXIOM Framework will attempt local execution. 305 | - A software suite is best represented as multiple, smaller tools instead of only one tool with dozens of commands. 306 | - Command names must be "tool-unique" across all YAML files within the `inventory folder`. 307 | - Concise but sufficiently-detailed command names, input names, and notes greatly improve the user experience. 308 | - Multi-line commands are not recommended as a substitute for writing real scripts in typical formats. 309 | 310 | ## Known Limitations 311 | 312 | - Doesn't set subprogram environment variables on its own 313 | - Doesn't do *anything* with non-STDOUT or non-PROMPT outputs 314 | - Doesn't track depth level for multiple interactive subprogram prompt changes 315 | - Doesn't clean up `bash` subprocesses after exiting interactive subprograms 316 | - Doesn't work well for subprograms that only exit upon receiving an interrupt 317 | - Doesn't work well for long-running, interactive subprograms with infrequent output 318 | - Fails to install PTF tools when the installation process requires user input 319 | - Fails to detect PTF tools not installed using organizational directories (not default) 320 | - Fails when running interactive "exit" commands in the wrong runtime context 321 | - Input placeholder text could hypothetically collide with some subprogram's syntax 322 | 323 | ## Feedback & Support 324 | 325 | Feel free to [open an issue](https://github.com/mikeiacovacci/axiom-framework/issues/new/choose) in any of the following 326 | scenarios: 327 | 328 | 1. [Bug](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=bug.md&labels=&title=%5BBUG%5D) 329 | 2. [Security weakness](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=security-weakness.md&labels=&title=%5BSECURITY%5D) 330 | not addressed [above](#security) 331 | 3. Significant [UI/UX problem](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=ui-ux-problem.md&labels=&title=%5BUI%2FUX%5D) 332 | 4. Missing or inaccurate [documentation](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=documentation.md&labels=&title=%5BDOCUMENTATION%5D) 333 | 5. [Inefficiency](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=inefficiency.md&labels=&title=%5BINEFFICIENCY%5D) 334 | (e.g. algorithmic) 335 | 6. Request for a reasonable [new feature](https://github.com/mikeiacovacci/axiom-framework/issues/new?assignees=&filename=feature-request.md&labels=&title=%5BFEATURE%5D) 336 | 337 | Please **do not** open any issues in the following scenarios: 338 | 339 | 1. Incorrect, inaccurate, or outdated toolkit data 340 | 2. Command execution failure due to toolkit data only 341 | 3. PTF issue unrelated to AXIOM Framework 342 | 4. A [known limitation](#known-limitations) (without proposing a solution) 343 | 5. Request for feature that broadly expands the project scope 344 | 6. Request for feature that harms end user security or privacy 345 | 7. Request for feature antithetical to the project ethos 346 | 347 | ## License 348 | 349 | AXIOM Framework is made available under the [Apache 2.0 License](https://www.apache.org/licenses/LICENSE-2.0). 350 | -------------------------------------------------------------------------------- /lib/classes.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Mike Iacovacci 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import lib.config as config 16 | from lib.config import print_error 17 | 18 | from os import devnull, path 19 | from pexpect import exceptions, pty_spawn 20 | from prompt_toolkit import prompt, PromptSession 21 | from prompt_toolkit.auto_suggest import AutoSuggestFromHistory 22 | from prompt_toolkit.history import FileHistory 23 | from queue import Queue 24 | from re import search 25 | from shlex import split 26 | from subprocess import call, PIPE, Popen, STDOUT 27 | from threading import Event 28 | from time import sleep 29 | 30 | 31 | class AxiomAction: 32 | """ A fully-completed, ready-to-execute tool command requiring no user input """ 33 | 34 | def __init__(self, name, prompt_type, execution_type, text, output_list, note): 35 | self.execution_type = execution_type 36 | self.name = name 37 | self.note = note 38 | self.output_list = output_list 39 | self.prompt_type = prompt_type 40 | self.text = text 41 | 42 | def cli_print(self): 43 | """ SUMMARY: displays the executable action text to the user, not stylized 44 | INPUT: none, reads values from self 45 | OUTPUT: none, only prints to the screen """ 46 | 47 | print() 48 | 49 | if isinstance(self.text, str): 50 | print(self.text) 51 | elif isinstance(self.text, list): 52 | line = 0 53 | while line < self.text.__len__(): 54 | print(self.text[line]) 55 | line += 1 56 | 57 | print() 58 | 59 | def confirm_and_execute(self, tool): 60 | """ SUMMARY: asks user to confirm execution of the command/action before proceeding 61 | INPUT: AxiomTool object 62 | OUTPUT: False if not confirmed, True if confirmed, after command/action executes """ 63 | 64 | self.show() 65 | response = input("\n[AXIOM] Execute? [Y/n] ") 66 | 67 | if response not in ["Y", "y", "Yes", "yes"]: 68 | return False 69 | else: 70 | self.run(tool) 71 | return True 72 | 73 | def existing_subprocess(self): 74 | """ SUMMARY: checks dispatch for existing subprocess with matching prompt type 75 | INPUT: none, reads values from self 76 | OUTPUT: True or False """ 77 | 78 | i = 0 79 | while i < dispatch.subprocesses.__len__(): 80 | if self.prompt_type == dispatch.subprocesses[i].current_prompt: 81 | return True 82 | i += 1 83 | 84 | return False 85 | 86 | def extract_ending_prompt(self): 87 | """ SUMMARY: determines the ending prompt of interactive command/action by processing output_list items 88 | INPUT: none, reads values from self 89 | OUTPUT: string containing prompt name, empty string if not found, or False if not interactive """ 90 | 91 | ending_prompt = str() 92 | 93 | if self.execution_type != "interactive": 94 | return False 95 | 96 | for x in self.output_list: 97 | if isinstance(x, tuple): 98 | if x[0] == "PROMPT": 99 | ending_prompt = x[1] 100 | break 101 | 102 | return ending_prompt 103 | 104 | def print_text(self): 105 | """ SUMMARY: displays the executable action text to the user, stylized 106 | INPUT: none, reads values from self 107 | OUTPUT: none, only prints to the screen """ 108 | 109 | if isinstance(self.text, str): 110 | print("\n TEXT: " + self.text) 111 | elif isinstance(self.text, list): 112 | print("\n TEXT: ", end="") 113 | print(self.text[0]) 114 | line = 1 115 | while line < self.text.__len__(): 116 | print(" " + self.text[line]) 117 | line += 1 118 | 119 | def run(self, tool): 120 | """ SUMMARY: checks if tool is compatible/installed and calls execution function for matching execution type 121 | INPUT: AxiomTool object 122 | OUTPUT: none """ 123 | 124 | if self.prompt_type == "bash" and not self.existing_subprocess(): 125 | 126 | if not tool.platform_matches(): 127 | print_error(str("\nERROR: Cannot execute " + tool.name + " (" + tool.platform + ") on " + 128 | config.axiom.platform)) 129 | dispatch.continue_trigger.set() 130 | return 131 | 132 | if tool.is_installed(): 133 | pass 134 | else: 135 | if tool.install(): 136 | self.show() 137 | print() 138 | else: 139 | if tool.proceed_despite_uninstalled(): 140 | pass 141 | else: 142 | dispatch.continue_trigger.set() 143 | return 144 | 145 | elif self.prompt_type != "other" and not self.existing_subprocess(): 146 | print_error("\nERROR: Prompt type incompatible with current runtime") 147 | dispatch.continue_trigger.set() 148 | return 149 | 150 | multiple_lines = False 151 | 152 | if isinstance(self, AxiomCommand): 153 | if isinstance(self.text[0], list): 154 | multiple_lines = True 155 | elif isinstance(self, AxiomAction): 156 | if isinstance(self.text, list): 157 | multiple_lines = True 158 | 159 | if self.execution_type == "standalone": 160 | if multiple_lines: 161 | self.run_multiline_standalone() 162 | else: 163 | self.run_standalone() 164 | elif self.execution_type == "autonomous": 165 | if multiple_lines: 166 | print_error("ERROR: Autonomous multi-line commands are unsupported") 167 | else: 168 | self.run_autonomous() 169 | elif self.execution_type == "interactive": 170 | self.run_interactive() 171 | elif self.execution_type == "NX": 172 | if multiple_lines: 173 | self.run_multiline_nx() 174 | else: 175 | self.run_nx() 176 | 177 | def run_autonomous(self): 178 | """ SUMMARY: executes autonomous action as subprocess (blocking) or queues action as a task (if interactive) 179 | INPUT: none, reads values from self 180 | OUTPUT: no return values """ 181 | 182 | if self.prompt_type == "bash" and not self.existing_subprocess(): 183 | try: 184 | print() 185 | call(self.text, shell=True) 186 | 187 | except OSError: 188 | print_error("ERROR: Failed to execute via call()") 189 | 190 | else: 191 | dispatch.tasking.put(AxiomInteractiveTask(self.text, self.prompt_type, self.prompt_type)) 192 | dispatch.monitor_task_queue() 193 | 194 | dispatch.continue_trigger.set() 195 | 196 | def run_interactive(self): 197 | """ SUMMARY: creates and queues an AxiomInteractiveTask object for execution 198 | INPUT: none, reads values from self 199 | OUTPUT: no return values """ 200 | 201 | ending_prompt = self.extract_ending_prompt() 202 | if ending_prompt is not False: 203 | dispatch.tasking.put(AxiomInteractiveTask(self.text, self.prompt_type, ending_prompt)) 204 | dispatch.monitor_task_queue() 205 | 206 | dispatch.continue_trigger.set() 207 | 208 | def run_multiline_nx(self): 209 | """ SUMMARY: prints multi-line action text to the screen 210 | INPUT: none, reads values from self 211 | OUTPUT: no return values, only prints to screen """ 212 | 213 | print() 214 | line = 0 215 | while line < self.text.__len__(): 216 | print(self.text[line]) 217 | line += 1 218 | dispatch.continue_trigger.set() 219 | 220 | def run_multiline_standalone(self): 221 | """ SUMMARY: executes multi-line action as subprocess or queues action execution as a task (if interactive) 222 | INPUT: none, reads values from self 223 | OUTPUT: none """ 224 | 225 | if self.prompt_type == "bash" and not self.existing_subprocess(): 226 | try: 227 | print() 228 | proc = Popen(["bash", "-i"], shell=True, stdin=PIPE, stdout=PIPE) 229 | 230 | i = 0 231 | while proc.returncode is None: 232 | if i < self.text.__len__(): 233 | proc.stdin.write(self.text[i].encode()) 234 | proc.stdin.write("\n".encode()) 235 | proc.stdin.flush() 236 | i += 1 237 | else: 238 | proc.stdin.close() 239 | 240 | proc.poll() 241 | 242 | except OSError: 243 | print_error("ERROR: Failed to execute via Popen()") 244 | 245 | else: 246 | dispatch.tasking.put(AxiomInteractiveTask(self.text, self.prompt_type, self.prompt_type)) 247 | dispatch.monitor_task_queue() 248 | 249 | dispatch.continue_trigger.set() 250 | 251 | def run_nx(self): 252 | """ SUMMARY: prints single-line action text to the screen 253 | INPUT: none, reads values from self 254 | OUTPUT: no return values, only prints to screen """ 255 | 256 | print() 257 | print(self.text) 258 | print() 259 | dispatch.continue_trigger.set() 260 | 261 | def run_standalone(self): 262 | """ SUMMARY: executes action as a subprocess (blocking) or queues action execution as a task (if interactive) 263 | INPUT: none, reads values from self 264 | OUTPUT: none """ 265 | 266 | if self.prompt_type == "bash" and not self.existing_subprocess(): 267 | try: 268 | print() 269 | call(split(self.text)) 270 | 271 | except OSError: 272 | print_error("ERROR: Failed to execute via call()") 273 | 274 | else: 275 | dispatch.tasking.put(AxiomInteractiveTask(self.text, self.prompt_type, self.prompt_type)) 276 | dispatch.monitor_task_queue() 277 | 278 | dispatch.continue_trigger.set() 279 | 280 | def show(self): 281 | """ SUMMARY: displays detailed information about the action to the user 282 | INPUT: none, reads values from self 283 | OUTPUT: none, only prints to the screen """ 284 | 285 | print("\n NAME: " + self.name + 286 | "\n TYPE: " + self.execution_type + " action (" + self.prompt_type + ")" 287 | "\n NOTE: " + self.note) 288 | 289 | self.print_text() 290 | 291 | 292 | class AxiomCommand(AxiomAction): 293 | """ The general syntax, including data-type placeholders, for an instruction to execute """ 294 | 295 | def __init__(self, name, prompt_type, execution_type, text, output_list, note, input_list): 296 | """ SUMMARY: creates AxiomCommand objects, inherits from AxiomAction class 297 | INPUT: multiples values at instantiation 298 | OUTPUT: none, instantiates AxiomCommand object """ 299 | 300 | super().__init__(name, prompt_type, execution_type, text, output_list, note) 301 | self.input_list = input_list 302 | 303 | def build(self): 304 | """ SUMMARY: interactively prompts user, possibly more than once, to enter/select all command input values 305 | INPUT: none, reads values from self 306 | OUTPUT: returns finalized command text, either a string or list of strings """ 307 | 308 | input_count = 0 309 | 310 | if isinstance(self.text[0], str): 311 | token_count = 0 312 | built_text = str() 313 | while token_count < self.text.__len__() or input_count < self.input_list.__len__(): 314 | if token_count < self.text.__len__(): 315 | built_text += self.text[token_count] 316 | token_count += 1 317 | if input_count < self.input_list.__len__(): 318 | built_text += self.input_build_prompt(input_count) 319 | input_count += 1 320 | else: 321 | built_text = [] 322 | current_line = 0 323 | while current_line < self.text.__len__(): 324 | line_tokens = self.text[current_line].__len__() 325 | current_token = 0 326 | line_inputs = line_tokens - 1 327 | current_input = 0 328 | built_line = str() 329 | while current_token < line_tokens or current_input < line_inputs: 330 | if current_token < line_tokens: 331 | built_line += self.text[current_line][current_token] 332 | current_token += 1 333 | if current_input < line_inputs: 334 | built_line += self.input_build_prompt(input_count) 335 | current_input += 1 336 | input_count += 1 337 | built_text.append(built_line) 338 | current_line += 1 339 | 340 | return built_text 341 | 342 | def build_with_placeholders(self): 343 | """ SUMMARY: creates command text containing placeholders for user preview before confirming execution 344 | INPUT: none, reads values from self 345 | OUTPUT: returns string or list of strings containing placeholders character sequences """ 346 | 347 | input_count = 0 348 | 349 | if isinstance(self.text[0], str): 350 | token_count = 0 351 | built_text = str() 352 | while token_count < self.text.__len__() or input_count < self.input_list.__len__(): 353 | if token_count < self.text.__len__(): 354 | built_text += self.text[token_count] 355 | token_count += 1 356 | if input_count < self.input_list.__len__(): 357 | built_text += str("{" + self.input_list[input_count][1] + "}") 358 | input_count += 1 359 | else: 360 | built_text = [] 361 | current_line = 0 362 | while current_line < self.text.__len__(): 363 | line_tokens = self.text[current_line].__len__() 364 | current_token = 0 365 | line_inputs = line_tokens - 1 366 | current_input = 0 367 | built_line = str() 368 | while current_token < line_tokens or current_input < line_inputs: 369 | if current_token < line_tokens: 370 | built_line += self.text[current_line][current_token] 371 | current_token += 1 372 | if current_input < line_inputs: 373 | built_line += str("{" + self.input_list[input_count][1] + "}") 374 | current_input += 1 375 | input_count += 1 376 | built_text.append(built_line) 377 | current_line += 1 378 | 379 | return built_text 380 | 381 | def cli_print(self): 382 | """ SUMMARY: prints command text to the screen (not stylized), overrides inherited AxiomAction function 383 | INPUT: none, reads values from self 384 | OUTPUT: none, only prints to the screen """ 385 | 386 | text = self.build() 387 | 388 | print() 389 | 390 | if isinstance(text, str): 391 | print(text) 392 | elif isinstance(text, list): 393 | line = 0 394 | while line < text.__len__(): 395 | print(text[line]) 396 | line += 1 397 | 398 | print() 399 | 400 | def input_build_prompt(self, input_count): 401 | """ SUMMARY: prompts user to enter, and auto-suggests, command inputs to replace placeholder values 402 | INPUT: current command input number (int), also reads values from self 403 | OUTPUT: returns a user-supplied or user-selected string value """ 404 | 405 | input_type = self.input_list[input_count][1] 406 | prompt_text = str("[AXIOM] Enter " + self.input_list[input_count][0] + ": ") 407 | 408 | if input_type in ["STRMENU", "INTMENU"]: 409 | option_name = self.input_list[input_count][0] 410 | option_list = self.input_list[input_count][2] 411 | response = self.option_prompt(option_name, option_list) 412 | return response 413 | elif input_type in ["STR", "INT", "IPV4", "IPV6", "IPV4RNGE", "IPV6RNGE", "IPV4CIDR", "IPV6CIDR", "MAC", "FILE", 414 | "RLATVPTH", "FULLPATH", "DOMAIN", "HTTPURL", "HTTPSURL", "WEBURL"]: 415 | 416 | if input_type == "HTTPSURL": 417 | history_file = str(config.axiom.history_folder + "/WEBURL" + ".axiom") 418 | else: 419 | history_file = str(config.axiom.history_folder + "/" + input_type + ".axiom") 420 | 421 | session = PromptSession(history=FileHistory(history_file)) 422 | response = session.prompt(prompt_text, auto_suggest=AutoSuggestFromHistory()) 423 | return response 424 | else: 425 | response = prompt(prompt_text) 426 | return response 427 | 428 | @staticmethod 429 | def option_prompt(option_name, option_list): 430 | """ SUMMARY: infinite loop prompting user to select a listed STRMENU or INTMENU option 431 | INPUT: option_name (str) and option_list (list) variables created from input_list values 432 | OUTPUT: string value from the option corresponding to the user's selection """ 433 | 434 | while True: 435 | print("\n" + option_name + "\n") 436 | 437 | count = 0 438 | while count < option_list.__len__(): 439 | print(" " + str(count + 1) + "\t" + str(option_list[count])) 440 | count += 1 441 | 442 | number = prompt("\n[AXIOM] Select an option: ") 443 | 444 | try: 445 | number = int(number) 446 | number -= 1 447 | except (ValueError, TypeError): 448 | number = -1 449 | 450 | if 0 <= number < option_list.__len__(): 451 | return option_list[number] 452 | 453 | def print_text(self): 454 | """ SUMMARY: prints command text to the screen (stylized), overrides inherited AxiomAction function 455 | INPUT: none, reads values from self 456 | OUTPUT: none, only prints to the screen """ 457 | 458 | text_with_placeholders = self.build_with_placeholders() 459 | if isinstance(text_with_placeholders, str): 460 | print("\n TEXT: " + text_with_placeholders) 461 | elif isinstance(text_with_placeholders, list): 462 | print("\n TEXT: ", end="") 463 | print(text_with_placeholders[0]) 464 | line = 1 465 | while line < text_with_placeholders.__len__(): 466 | print(" " + text_with_placeholders[line]) 467 | line += 1 468 | 469 | def run_autonomous(self): 470 | """ SUMMARY: builds and runs command as subprocess (blocking) or queues task for interactive execution 471 | overrides inherited AxiomAction function 472 | INPUT: none, reads values from self 473 | OUTPUT: none """ 474 | 475 | text = self.build() 476 | if self.prompt_type == "bash" and not self.existing_subprocess(): 477 | try: 478 | print() 479 | call(text, shell=True) 480 | 481 | except OSError: 482 | print_error("ERROR: Failed to execute via call()") 483 | 484 | else: 485 | dispatch.tasking.put(AxiomInteractiveTask(text, self.prompt_type, self.prompt_type)) 486 | dispatch.monitor_task_queue() 487 | 488 | dispatch.continue_trigger.set() 489 | 490 | def run_interactive(self): 491 | """ SUMMARY: builds command text and builds/queues interactive execution task 492 | overrides inherited AxiomAction function 493 | INPUT: none, reads values from self 494 | OUTPUT: none """ 495 | 496 | text = self.build() 497 | ending_prompt = self.extract_ending_prompt() 498 | if ending_prompt is not False: 499 | dispatch.tasking.put(AxiomInteractiveTask(text, self.prompt_type, ending_prompt)) 500 | dispatch.monitor_task_queue() 501 | 502 | dispatch.continue_trigger.set() 503 | 504 | def run_multiline_nx(self): 505 | """ SUMMARY: builds and prints multi-line command text to screen, overrides inherited AxiomAction function 506 | INPUT: none, reads values from self 507 | OUTPUT: no return values, only prints to screen """ 508 | 509 | text = self.build() 510 | print() 511 | line = 0 512 | while line < self.text.__len__(): 513 | print(text[line]) 514 | line += 1 515 | dispatch.continue_trigger.set() 516 | 517 | def run_multiline_standalone(self): 518 | """ SUMMARY: builds and executes command as subprocess or queues task for interactive execution 519 | overrides inherited AxiomAction function 520 | INPUT: none, reads values from self 521 | OUTPUT: no return values """ 522 | 523 | text = self.build() 524 | if self.prompt_type == "bash" and not self.existing_subprocess(): 525 | try: 526 | print() 527 | proc = Popen(["bash", "-i"], shell=True, stdin=PIPE, stdout=PIPE) 528 | 529 | i = 0 530 | while proc.returncode is None: 531 | if i < text.__len__(): 532 | proc.stdin.write(text[i].encode()) 533 | proc.stdin.write("\n".encode()) 534 | proc.stdin.flush() 535 | i += 1 536 | else: 537 | proc.stdin.close() 538 | 539 | proc.poll() 540 | 541 | except OSError: 542 | print_error("ERROR: Failed to execute via Popen()") 543 | else: 544 | dispatch.tasking.put(AxiomInteractiveTask(text, self.prompt_type, self.prompt_type)) 545 | dispatch.monitor_task_queue() 546 | 547 | dispatch.continue_trigger.set() 548 | 549 | def run_nx(self): 550 | """ SUMMARY: builds and displays command text to screen, overrides inherited AxiomAction function 551 | INPUT: none, reads values from self 552 | OUTPUT: no return values, only prints to the screen """ 553 | 554 | text = self.build() 555 | print() 556 | print(text) 557 | print() 558 | dispatch.continue_trigger.set() 559 | 560 | def run_standalone(self): 561 | """ SUMMARY: builds and executes command as subprocess (blocking) or queues interactive task for execution 562 | overrides inherited AxiomAction function 563 | INPUT: none, reads values from self 564 | OUTPUT: no return values """ 565 | 566 | text = self.build() 567 | if self.prompt_type == "bash" and not self.existing_subprocess(): 568 | try: 569 | print() 570 | call(split(text)) 571 | 572 | except OSError: 573 | print_error("ERROR: Failed to execute via call()") 574 | else: 575 | dispatch.tasking.put(AxiomInteractiveTask(text, self.prompt_type, self.prompt_type)) 576 | dispatch.monitor_task_queue() 577 | 578 | dispatch.continue_trigger.set() 579 | 580 | def show(self): 581 | """ SUMMARY: displays detailed information about the command, overrides inherited AxiomAction function 582 | INPUT: none, reads values from self 583 | OUTPUT: none, only prints to the screen """ 584 | 585 | print("\n NAME: " + self.name + 586 | "\n TYPE: " + self.execution_type + " command (" + self.prompt_type + ")" 587 | "\n NOTE: " + self.note) 588 | 589 | self.print_text() 590 | 591 | 592 | class AxiomDispatcher: 593 | """ creates, manages, and interacts with subprocesses that require interactive input """ 594 | 595 | def __init__(self): 596 | self.continue_trigger = Event() 597 | self.subprocesses = [] 598 | self.tasking = Queue(maxsize=0) 599 | self.trigger = Event() 600 | 601 | def check_for_ambiguous_target(self, current_task): 602 | """ SUMMARY: detects existing subprocesses with prompt types that match a task's ending prompt type 603 | INPUT: current_task, an AxiomInteractiveTask object from the "tasking" queue 604 | OUTPUT: True or False """ 605 | 606 | prompt_type = current_task.ending_prompt 607 | 608 | for x in self.subprocesses: 609 | if x.current_prompt == prompt_type: 610 | return True 611 | 612 | return False 613 | 614 | @staticmethod 615 | def get_subprocess_output_detect_prompt(proc, pattern): 616 | """ SUMMARY: prints subprocess output to the screen while searching for an interactive prompt 617 | INPUT: 1) a pseudoterminal subprocess object from pty_spawn and 2) a regex prompt pattern (str) 618 | OUTPUT: no return values, only prints to the screen """ 619 | 620 | timeout = 0 621 | safety_timer = 0 622 | 623 | while True: 624 | try: 625 | print(proc.readline().decode(), end='') 626 | except exceptions.TIMEOUT: 627 | if search(pattern, proc.before.decode()): 628 | if timeout >= config.axiom.pattern_timeout: 629 | print(proc.before.decode()) 630 | break 631 | else: 632 | timeout += 1 633 | sleep(1) 634 | continue 635 | else: 636 | safety_timer += 1 637 | sleep(1) 638 | if safety_timer >= config.axiom.safety_timeout: 639 | proc.sendline() 640 | continue 641 | else: 642 | timeout = 0 643 | safety_timer = 0 644 | 645 | def handle_new_tasks(self): 646 | """ SUMMARY: gets AxiomInteractiveTask objects from queue and routes tasks based on runtime context 647 | INPUT: self, gets objects from "tasking" queue 648 | OUTPUT: no return value, returns when task is routed """ 649 | 650 | if not self.tasking.empty(): 651 | current_task = self.tasking.get() 652 | if self.matching_subprocess(current_task) >= 0: 653 | target = self.matching_subprocess(current_task) 654 | if current_task.prompt_change: 655 | if self.check_for_ambiguous_target(current_task): 656 | print_error("\nERROR: Cannot create subprocess with same prompt type as existing subprocess") 657 | self.tasking.task_done() 658 | return 659 | self.read_and_transmit(target, current_task) 660 | self.tasking.task_done() 661 | return 662 | elif current_task.starting_prompt == "bash": 663 | if self.check_for_ambiguous_target(current_task): 664 | print_error("\nERROR: Cannot create subprocess with same prompt type as existing subprocess") 665 | self.tasking.task_done() 666 | return 667 | self.spawn_and_transmit(current_task) 668 | self.tasking.task_done() 669 | return 670 | else: 671 | print_error("\nERROR: Prompt type incompatible with current runtime") 672 | self.tasking.task_done() 673 | return 674 | 675 | def matching_subprocess(self, current_task): 676 | """ SUMMARY: locates existing subprocess with identical prompt type to queued task 677 | INPUT: current_task, an AxiomInteractiveTask object from the "tasking" queue 678 | OUTPUT: integer, zero or positive if match found, -1 if no match """ 679 | 680 | i = 0 681 | while i < self.subprocesses.__len__(): 682 | if current_task.starting_prompt == self.subprocesses[i].current_prompt: 683 | return i 684 | else: 685 | i += 1 686 | 687 | return -1 688 | 689 | def monitor_task_queue(self): 690 | """ calls any required functions related to new tasks in the queue """ 691 | 692 | self.handle_new_tasks() 693 | 694 | def read_and_transmit(self, target, current_task): 695 | """ SUMMARY: prints prior program output, transmits text to existing subprocess, and updates the prompt 696 | INPUT: targeted subprocess number (INT) and AxiomInteractiveTask object from "tasking" queue 697 | OUTPUT: no return values """ 698 | 699 | proc = self.subprocesses[target].process 700 | 701 | while True: 702 | try: 703 | print(proc.readline().decode(), end='') 704 | except exceptions.TIMEOUT: 705 | break 706 | 707 | self.transmit_text(current_task, proc) 708 | 709 | self.subprocesses[target].current_prompt = current_task.ending_prompt 710 | self.subprocesses[target].prompt_pattern = current_task.ending_prompt_pattern 711 | dispatch.continue_trigger.set() 712 | 713 | def spawn_and_transmit(self, current_task): 714 | """ SUMMARY: creates a new subprocess, transmits a command's/action's executable text, and updates the prompt 715 | INPUT: an AxiomInteractiveTask object from the "tasking" queue 716 | OUTPUT: no return values """ 717 | 718 | try: 719 | self.subprocesses.append(AxiomExecutingSubprocess(current_task.starting_prompt, 720 | pty_spawn.spawn("/bin/bash -i", 721 | timeout=config.axiom.pty_timeout))) 722 | 723 | except OSError: 724 | print_error("ERROR: Failed to spawn /bin/bash subprocess") 725 | exit(1) 726 | 727 | else: 728 | target = self.matching_subprocess(current_task) 729 | proc = self.subprocesses[target].process 730 | 731 | self.transmit_text(current_task, proc) 732 | 733 | self.subprocesses[target].current_prompt = current_task.ending_prompt 734 | self.subprocesses[target].prompt_pattern = current_task.ending_prompt_pattern 735 | dispatch.continue_trigger.set() 736 | 737 | def transmit_text(self, current_task, proc): 738 | """ SUMMARY: transmits line-buffered input to a subprocess and waits for & displays the subprocess's output 739 | INPUT: 1) an AxiomInteractiveTask object and 2) a pseudoterminal subprocess object from pty_spawn 740 | OUTPUT: no return values, only prints to the screen """ 741 | 742 | pattern = str(current_task.ending_prompt_pattern + "$") 743 | 744 | try: 745 | if isinstance(current_task.text, str): 746 | proc.sendline(current_task.text) 747 | elif isinstance(current_task.text, list): 748 | i = 0 749 | while i < current_task.text.__len__(): 750 | proc.sendline(current_task.text[i]) 751 | i += 1 752 | 753 | except OSError: 754 | print_error("ERROR: Failed to transmit command") 755 | exit(1) 756 | 757 | else: 758 | self.get_subprocess_output_detect_prompt(proc, pattern) 759 | 760 | 761 | class AxiomExecutingSubprocess: 762 | """ structure for managing subprocesses that require interactive input """ 763 | 764 | def __init__(self, current_prompt, process): 765 | self.current_prompt = current_prompt 766 | self.process = process 767 | self.prompt_pattern = None 768 | 769 | 770 | class AxiomInteractiveTask: 771 | """ defines tasks sent to AxiomDispatcher queue for working with interactive subprocesses """ 772 | 773 | def __init__(self, text, starting_prompt, ending_prompt): 774 | """ SUMMARY: creates object, to be queued, for handling interactive execution tasks 775 | INPUT: the finalized command/action text (str or list), and starting + ending prompt type names 776 | OUTPUT: self, instantiates an AxiomInteractiveTask object """ 777 | 778 | self.ending_prompt = ending_prompt 779 | self.starting_prompt = starting_prompt 780 | self.text = text 781 | 782 | self.prompt_change = self.detect_prompt_change() 783 | 784 | self.ending_prompt_pattern = self.resolve_ending_prompt_pattern() 785 | 786 | def detect_prompt_change(self): 787 | """ SUMMARY: compares two prompt type names, called by AxiomInteractiveTask init method 788 | INPUT: self, two string values that represent prompt type names 789 | OUTPUT: True or False based on string comparison """ 790 | 791 | if self.starting_prompt == self.ending_prompt: 792 | return False 793 | else: 794 | return True 795 | 796 | def resolve_ending_prompt_pattern(self): 797 | """ SUMMARY: extracts ending prompt pattern from global config object 798 | INPUT: self and global config object 799 | OUTPUT: string containing the appropriate prompt pattern """ 800 | 801 | if self.prompt_change: 802 | for x in config.axiom.prompts: 803 | if x[0] == self.ending_prompt: 804 | return x[1] 805 | else: 806 | for x in config.axiom.prompts: 807 | if x[0] == self.starting_prompt: 808 | return x[1] 809 | 810 | 811 | class AxiomToolkit: 812 | """ A collection of related tools """ 813 | 814 | def __init__(self, name, location, tool_name_list): 815 | self.location = location 816 | self.name = name 817 | self.tool_name_list = tool_name_list 818 | 819 | 820 | class AxiomTool: 821 | """ an executable program with related commands and actions """ 822 | 823 | def __init__(self, name, platform, ptf_module, description, action_list, command_list): 824 | self.action_list = action_list 825 | self.combined_list = [] 826 | self.command_list = command_list 827 | self.description = description 828 | self.name = name 829 | self.platform = platform 830 | self.ptf_module = ptf_module 831 | 832 | def initialize_combined_list(self): 833 | """ SUMMARY: creates alphabetically-ordered list of command/action names 834 | INPUT: self, reads action_list and command_list variables 835 | OUTPUT: none, modifies combined_list variable """ 836 | 837 | self.combined_list = [] 838 | x = 0 839 | while x < self.action_list.__len__(): 840 | self.combined_list.append(self.action_list[x].name) 841 | x += 1 842 | y = 0 843 | while y < self.command_list.__len__(): 844 | self.combined_list.append(self.command_list[y].name) 845 | y += 1 846 | 847 | self.combined_list = sorted(self.combined_list, key=str.casefold) 848 | 849 | def install(self): 850 | """ SUMMARY: prompts user and installs undetected tools to local system via PTF when possible 851 | INPUT: none, reads values from self & conditionally prompts user for interactive input 852 | OUTPUT: True or False """ 853 | 854 | if self.ptf_module not in ["", None]: 855 | answer = input("[AXIOM] Install " + self.name + " via PTF? [Y/n] ") 856 | if answer not in ["Y", "y", "Yes", "yes"]: 857 | return False 858 | else: 859 | if config.axiom.platform.lower() != "linux": 860 | print_error(str("ERROR: Unable to run PTF on " + config.axiom.platform)) 861 | return False 862 | else: 863 | input_text = str("python3 ./ptf --no-network-connection << EOF\n" + 864 | str("use " + self.ptf_module + "\n") + 865 | "install\n" + 866 | "EOF\n") 867 | try: 868 | call(input_text, shell=True, cwd=config.axiom.ptf_folder) 869 | return True 870 | 871 | except OSError: 872 | print_error("ERROR: Failed to execute PTF") 873 | exit(1) 874 | else: 875 | return False 876 | 877 | def is_installed(self): 878 | """ SUMMARY: checks local system for installed tool via 1) PTF and 2) 'which' command 879 | INPUT: none. reads values from self 880 | OUTPUT: True or False """ 881 | 882 | ptf_config_file = str(config.axiom.ptf_folder + "/config/ptf.config") 883 | 884 | if self.ptf_module not in ["", None]: 885 | tool_module_file = str(config.axiom.ptf_folder + "/" + self.ptf_module + ".py") 886 | 887 | try: 888 | with open(ptf_config_file) as ptf_config: 889 | for line in enumerate(ptf_config): 890 | if search("^BASE_INSTALL_PATH=", line[1]): 891 | install_path = line[1].split("\"")[1] 892 | break 893 | 894 | except OSError: 895 | print_error(str("ERROR: Failed to extract PTF base install path from " + ptf_config_file)) 896 | exit(1) 897 | 898 | else: 899 | try: 900 | with open(tool_module_file) as module_file: 901 | for line in enumerate(module_file): 902 | if search("^INSTALL_LOCATION=", line[1]): 903 | location = line[1].split("\"")[1] 904 | break 905 | 906 | except OSError: 907 | print_error(str("ERROR: Failed to extract PTF install location from " + tool_module_file)) 908 | exit(1) 909 | 910 | else: 911 | folder = str(self.ptf_module.split("/")[1]) 912 | ptf_tool_folder = str(install_path + "/" + folder + "/" + location) 913 | 914 | if path.exists(ptf_tool_folder): 915 | return True 916 | else: 917 | return False 918 | 919 | text = str("which \"" + self.name + "\"") 920 | 921 | try: 922 | dev_null = open(devnull, 'w') 923 | if call(split(text), stdout=dev_null, stderr=STDOUT) == 0: 924 | return True 925 | else: 926 | return False 927 | 928 | except OSError: 929 | print_error(str("ERROR: Failed to run command " + text)) 930 | exit(1) 931 | 932 | def platform_matches(self): 933 | """ SUMMARY: compares tool platform against local platform value in global config object 934 | INPUT: none, reads values from self and config 935 | OUTPUT: True or False """ 936 | 937 | if self.platform.lower() == config.axiom.platform.lower(): 938 | return True 939 | else: 940 | return False 941 | 942 | def proceed_despite_uninstalled(self): 943 | """ SUMMARY: prompts user to confirm it's okay to execute the tool regardless of if it was detected 944 | INPUT: an AxiomTool object 945 | OUTPUT: True of False """ 946 | 947 | answer = input("[AXIOM] Unable to confirm " + self.name + " is installed. Proceed anyway? [Y/n] ") 948 | if answer not in ["Y", "y", "Yes", "yes"]: 949 | return False 950 | else: 951 | return True 952 | 953 | def resolve_command(self, number): 954 | """ SUMMARY: determines the object's type (command or action) and finds its ID value 955 | INPUT: command/action ID number integer 956 | OUTPUT: two-item tuple containing 1) "command", "action", or None and 2) ID value, -1 if unresolved """ 957 | 958 | if number >= 0 and number in range(self.combined_list.__len__()): 959 | command_name = self.combined_list[number] 960 | return self.resolve_command_name(command_name) 961 | else: 962 | return None, int(-1) 963 | 964 | def resolve_command_name(self, command_name): 965 | """ SUMMARY: finds the ID value of the supplied command/action name 966 | INPUT: command/action name string 967 | OUTPUT: tuple containing string ("command" or "action") and ID value (int), -1 if not found """ 968 | 969 | command_type = str() 970 | id_value = int(-1) 971 | 972 | x = 0 973 | action_count = self.action_list.__len__() 974 | while x < action_count: 975 | if self.action_list[x].name == command_name: 976 | command_type = "action" 977 | id_value = x 978 | x += 1 979 | 980 | y = 0 981 | command_count = self.command_list.__len__() 982 | while y < command_count: 983 | if self.command_list[y].name == command_name: 984 | command_type = "command" 985 | id_value = y 986 | y += 1 987 | 988 | return command_type, id_value 989 | 990 | def show(self): 991 | """ SUMMARY: displays tool information on the screen for the user 992 | INPUT: self, reads name, ptf_module, description, and combined_list variables 993 | OUTPUT: none, only prints to the screen """ 994 | 995 | print("\n NAME: " + str(self.name) + " (" + str(self.platform) + ")") 996 | 997 | if isinstance(self.ptf_module, str): 998 | print(" TOOL: " + str(self.ptf_module)) 999 | 1000 | print(" NOTE: " + str(self.description)) 1001 | 1002 | print("\nCommands\n") 1003 | i = 0 1004 | while i < self.combined_list.__len__(): 1005 | print(" " + str(i + 1) + "\t" + self.combined_list[i]) 1006 | i += 1 1007 | 1008 | 1009 | dispatch = AxiomDispatcher() 1010 | -------------------------------------------------------------------------------- /lib/functions.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Mike Iacovacci 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from lib.classes import * 16 | 17 | from colorama import Fore, Style 18 | from io import BytesIO 19 | from os import geteuid, listdir, mkdir, path, rename, remove 20 | from pickle import dump, load, PickleError 21 | from prompt_toolkit import prompt 22 | from prompt_toolkit.completion import FuzzyCompleter, WordCompleter 23 | from prompt_toolkit.styles import Style as ptkStyle 24 | from re import split 25 | from shutil import rmtree 26 | from sys import argv 27 | from requests import get, RequestException 28 | from yaml import safe_load_all, parser, scanner 29 | from zipfile import BadZipFile, LargeZipFile, ZipFile 30 | 31 | 32 | def able_to_merge(current_tool, tool_id, tools): 33 | """ SUMMARY: determines if merging YAML data from 2 files is possible without data loss 34 | INPUT: 1) an AxiomTool object, a tool ID value (int), and 2) a list of AxiomTool objects 35 | OUTPUT: True or False """ 36 | 37 | if tool_id < 0: 38 | return False 39 | 40 | match = tools[tool_id] 41 | temp_list = [] 42 | new_list = [] 43 | 44 | if current_tool[0]["ptf_module"] != match.ptf_module: 45 | return False 46 | if current_tool[0]["description"] != match.description: 47 | return False 48 | 49 | action_count = 0 50 | while action_count < match.action_list.__len__(): 51 | temp_list.append(match.action_list[action_count].name) 52 | action_count += 1 53 | 54 | command_count = 0 55 | while command_count < match.command_list.__len__(): 56 | temp_list.append(match.command_list[command_count].name) 57 | command_count += 1 58 | 59 | new_count = 0 60 | while new_count < current_tool[1]["commands"].__len__(): 61 | new_list.append(list(current_tool[1]["commands"][new_count].keys())[0]) 62 | new_count += 1 63 | 64 | for command_name in new_list: 65 | if command_name in temp_list: 66 | return False 67 | 68 | return True 69 | 70 | 71 | def axiom_help(): 72 | """ SUMMARY: displays helpful CLI usage details with examples 73 | INPUT: none 74 | OUTPUT: none, only prints to the screen """ 75 | 76 | print("\n" + "Standard usage: ./axiom [MODE] [TOOL] [NUM]" 77 | "\n" + "" + 78 | "\n" + " ./axiom show nmap" + 79 | "\n" + " ./axiom show mimikatz 1" + 80 | "\n" + " ./axiom build powershell 4" + 81 | "\n" + " ./axiom run hashcat 3" + 82 | "\n" + "" + 83 | "\n" + "Configuration management: ./axiom [MODE] [URL]" + 84 | "\n" + "" + 85 | "\n" + " ./axiom new" + 86 | "\n" + " ./axiom reload" + 87 | "\n" + " ./axiom init" + 88 | "\n" + " ./axiom init https://example.com/config.yml" + 89 | "\n") 90 | 91 | 92 | def axiom_prompt(tool_list, tool_names, tools): 93 | """ SUMMARY: main interactive prompt loop of the program, handles multiple tool selection loops 94 | INPUT: 1) list of two-item tuples (name, platform), 2) set of tool names, and 3) list of AxiomTool objects 95 | OUTPUT: exit value (int) to be immediately passed to exit() in __main__ """ 96 | 97 | exit_code = 1 98 | 99 | while exit_code > 0: 100 | exit_code = tool_selection_prompt(tool_list, tool_names, tools) 101 | 102 | return exit_code 103 | 104 | 105 | def branch(settings, tool_list, tools): 106 | """ SUMMARY: changes program flow based on user-supplied settings 107 | INPUT: 1) a three-item dictionary, 2) a de-duplicated list of tuples, and 3) list of AxiomTool objects 108 | OUTPUT: no return value, may exit the entire program """ 109 | 110 | if settings.get("mode") in [None, "reload", "init"]: 111 | return 112 | 113 | if settings.get("mode") == "new": 114 | new_generate_command() 115 | 116 | if settings.get("num") == -1: 117 | print_error("ERROR: Invalid command ID") 118 | exit(1) 119 | 120 | text = settings.get("tool") 121 | tool_id = disambiguate_tool_name(text, tool_list, tools) 122 | 123 | if tool_id < 0: 124 | print_error("ERROR: Invalid tool") 125 | exit(1) 126 | 127 | tool = tools[tool_id] 128 | 129 | if settings.get("mode") == "show": 130 | if settings.get("num") is None: 131 | tool.show() 132 | print() 133 | exit(0) 134 | elif int(settings.get("num") - 1) not in range(tool.combined_list.__len__()): 135 | print_error("ERROR: Invalid command specified") 136 | exit(1) 137 | else: 138 | number = int(settings.get("num") - 1) 139 | command_type, id_value = tool.resolve_command(number) 140 | if command_type == "action": 141 | tool.action_list[id_value].show() 142 | print() 143 | elif command_type == "command": 144 | tool.command_list[id_value].show() 145 | print() 146 | exit(0) 147 | 148 | if settings.get("mode") == "run": 149 | if settings.get("num") is None: 150 | print_error("ERROR: No command specified") 151 | exit(1) 152 | elif int(settings.get("num") - 1) not in range(tool.combined_list.__len__()): 153 | print_error("ERROR: Invalid command specified") 154 | exit(1) 155 | else: 156 | number = int(settings.get("num") - 1) 157 | command_type, id_value = tool.resolve_command(number) 158 | if command_type == "action": 159 | if tool.action_list[id_value].execution_type in ["standalone", "autonomous", "NX"]: 160 | tool.action_list[id_value].run(tool) 161 | else: 162 | print_error("ERROR: Selected action must be executed via interactive AXIOM prompt") 163 | exit(1) 164 | elif command_type == "command": 165 | if tool.command_list[id_value].execution_type in ["standalone", "autonomous", "NX"]: 166 | tool.command_list[id_value].run(tool) 167 | else: 168 | print_error("ERROR: Selected command must be executed via interactive AXIOM prompt") 169 | exit(1) 170 | exit(0) 171 | 172 | if settings.get("mode") == "build": 173 | if settings.get("num") is None: 174 | print_error("ERROR: No command specified") 175 | exit(1) 176 | elif int(settings.get("num") - 1) not in range(tool.combined_list.__len__()): 177 | print_error("ERROR: Invalid command specified") 178 | exit(1) 179 | else: 180 | number = int(settings.get("num") - 1) 181 | command_type, id_value = tool.resolve_command(number) 182 | if command_type == "action": 183 | tool.action_list[id_value].cli_print() 184 | elif command_type == "command": 185 | tool.command_list[id_value].cli_print() 186 | exit(0) 187 | 188 | 189 | def command_selection_prompt(tool): 190 | """ SUMMARY: prompts user to select a listed command/action for the current tool and calls the execution function 191 | INPUT: an AxiomTool object 192 | OUTPUT: none """ 193 | 194 | while True: 195 | tool.show() 196 | number = prompt('\n[AXIOM] Select command: ') 197 | 198 | if number == "back": 199 | return 200 | if number == "exit" or number == "quit": 201 | print("Exiting...") 202 | exit(0) 203 | 204 | if number == "": 205 | continue 206 | 207 | try: 208 | number = int(number) 209 | number -= 1 210 | except (ValueError, TypeError): 211 | number = -1 212 | 213 | if number not in range(tool.combined_list.__len__()): 214 | print_error("\nERROR: Invalid command specified") 215 | else: 216 | command_type, id_value = tool.resolve_command(number) 217 | 218 | if command_type == "action": 219 | confirmed = tool.action_list[id_value].confirm_and_execute(tool) 220 | elif command_type == "command": 221 | confirmed = tool.command_list[id_value].confirm_and_execute(tool) 222 | else: 223 | confirmed = False 224 | 225 | if confirmed: 226 | dispatch.continue_trigger.wait(timeout=None) 227 | dispatch.continue_trigger.clear() 228 | print() 229 | input("[AXIOM] Press ENTER to continue ") 230 | 231 | 232 | def create_missing_folder(folder): 233 | """ SUMMARY: checks if specified folder exists, creates it if it does not exist 234 | INPUT: a string specifying a folder on the filesystem 235 | OUTPUT: none, creates necessary folder if it does not exist """ 236 | 237 | if path.exists(folder): 238 | return 239 | else: 240 | try: 241 | mkdir(folder) 242 | except OSError: 243 | print_error(str("ERROR: Cannot create folder \"" + folder + "\"")) 244 | exit(1) 245 | 246 | 247 | def delete_and_recreate_folder(folder): 248 | """ SUMMARY: deletes specified folder, if it exists, and (re)creates it on the filesystem 249 | INPUT: a string specifying a folder on the filesystem 250 | OUTPUT: none, deletes files from the filesystem and/or creates necessary folders """ 251 | 252 | if path.exists(folder): 253 | try: 254 | rmtree(folder) 255 | except OSError: 256 | print_error(str("ERROR: Cannot delete folder \"" + folder + "\"")) 257 | exit(1) 258 | 259 | create_missing_folder(folder) 260 | 261 | 262 | def disambiguate_tool_name(text, tool_list, tools): 263 | """ SUMMARY: finds the user-intended tool ID for multi-platform tool names prompting the user as needed 264 | INPUT: 1) supplied tool name (str), 2) de-duplicated list of tuples, and 3) list of AxiomTool objects 265 | OUTPUT: tool ID value (int) or -1 if invalid number of platforms or no matching tool found """ 266 | 267 | platform_list = [] 268 | for x in tool_list: 269 | if x[0] == text: 270 | platform_list.append(x[1]) 271 | 272 | platform_list = sorted(platform_list, key=str.casefold) 273 | 274 | potential_tool = [] 275 | if platform_list.__len__() == 0: 276 | return -1 277 | elif platform_list.__len__() == 1: 278 | potential_tool.append(text) 279 | potential_tool.append(platform_list[0]) 280 | else: 281 | selection = 0 282 | while selection == 0: 283 | print("\nPlatforms\n") 284 | i = 0 285 | while i < platform_list.__len__(): 286 | print(" " + str(i + 1) + "\t" + platform_list[i]) 287 | i += 1 288 | platform = prompt('\n[AXIOM] Select platform: ') 289 | 290 | try: 291 | number = int(platform) 292 | except (ValueError, TypeError): 293 | number = 0 294 | 295 | if number > 0: 296 | if number <= platform_list.__len__(): 297 | potential_tool.append(text) 298 | potential_tool.append(platform_list[number - 1]) 299 | selection = 1 300 | 301 | return resolve_tool_id(potential_tool, tools) 302 | 303 | 304 | def download_and_extract_zip(zip_url, extracted_folder, destination_folder, human_name): 305 | """ SUMMARY: prepares the filesystem, downloads a ZIP file, and extracts it to a folder with a specified name 306 | INPUT: ZIP file URL, temporary folder name, destination folder name, and human-friendly name (all strings) 307 | OUTPUT: no return values, modifies the filesystem """ 308 | 309 | if path.exists(destination_folder): 310 | try: 311 | rmtree(destination_folder) 312 | except OSError: 313 | print_error(str("ERROR: Cannot prepare extraction location \"" + destination_folder + "\"")) 314 | exit(1) 315 | 316 | print("Downloading " + human_name + "...") 317 | 318 | try: 319 | request = get(zip_url) 320 | 321 | except RequestException: 322 | print_error(str("ERROR: Cannot download \"" + human_name + "\" from " + zip_url)) 323 | exit(1) 324 | 325 | else: 326 | if request.status_code == 200: 327 | 328 | try: 329 | zipfile = ZipFile(BytesIO(request.content)) 330 | zipfile.extractall(".") 331 | rename(extracted_folder, destination_folder) 332 | 333 | except (BadZipFile, LargeZipFile, OSError): 334 | print_error(str("ERROR: Cannot extract \"" + extracted_folder + "\"")) 335 | exit(1) 336 | 337 | else: 338 | print_error(str("ERROR: Failed to download \"" + human_name + "\"")) 339 | exit(1) 340 | 341 | 342 | def get_args(): 343 | """ SUMMARY: processes command-line arguments to modify overall program execution flow 344 | INPUT: none, checks argv for arguments supplied via CLI 345 | OUTPUT: three-item dictionary containing the mode type, tool name, and command/action number """ 346 | 347 | if argv.__len__() < 2: 348 | return {"mode": None, "tool": None, "num": None} 349 | 350 | elif argv.__len__() > 4: 351 | axiom_help() 352 | exit(1) 353 | 354 | elif argv.__len__() == 2: 355 | 356 | if argv[1] == "init": 357 | return {"mode": "init", "tool": None, "num": None} 358 | if argv[1] == "reload": 359 | return {"mode": "reload", "tool": None, "num": None} 360 | if argv[1] in ["n", "ne", "new", "-n", "--new"]: 361 | return {"mode": "new", "tool": None, "num": None} 362 | else: 363 | axiom_help() 364 | exit(1) 365 | 366 | elif argv.__len__() == 3 or 4: 367 | 368 | if argv[1] == "init": 369 | return {"mode": "init", "tool": str(argv[2]), "num": None} 370 | if argv[1] in ["s", "sh", "sho", "show", "-s", "--show"]: 371 | if argv.__len__() == 3: 372 | return {"mode": "show", "tool": str(argv[2]), "num": None} 373 | if argv.__len__() == 4: 374 | try: 375 | number = int(argv[3]) 376 | except (ValueError, TypeError): 377 | number = -1 378 | return {"mode": "show", "tool": str(argv[2]), "num": number} 379 | 380 | if argv[1] in ["r", "ru", "run", "-r", "--run"]: 381 | if argv.__len__() == 3: 382 | return {"mode": "run", "tool": str(argv[2]), "num": None} 383 | if argv.__len__() == 4: 384 | try: 385 | number = int(argv[3]) 386 | except (ValueError, TypeError): 387 | number = -1 388 | return {"mode": "run", "tool": str(argv[2]), "num": number} 389 | 390 | if argv[1] in ["b", "bu", "bui", "buil", "build", "-b", "--build"]: 391 | if argv.__len__() == 3: 392 | return {"mode": "build", "tool": str(argv[2]), "num": None} 393 | if argv.__len__() == 4: 394 | try: 395 | number = int(argv[3]) 396 | except (ValueError, TypeError): 397 | number = -1 398 | return {"mode": "build", "tool": str(argv[2]), "num": number} 399 | 400 | else: 401 | axiom_help() 402 | exit(1) 403 | 404 | else: 405 | axiom_help() 406 | exit(1) 407 | 408 | 409 | def get_input_types(input_types_list, text): 410 | """ SUMMARY: parses placeholder text to determine the type of input required for command/action execution 411 | INPUT: 1) list of all possible input types (strings), and 2) the command text (list or str) 412 | OUTPUT: a list of strings """ 413 | 414 | if isinstance(text, list): 415 | temporary_string = "" 416 | line_count = 0 417 | while line_count < text.__len__(): 418 | temporary_string += text[line_count] 419 | line_count += 1 420 | text = temporary_string 421 | 422 | used_input_types = [] 423 | end = text.__len__() 424 | indices = [i for i in range(end) if text.startswith("{", i)] 425 | hit_count = 0 426 | 427 | while hit_count < indices.__len__(): 428 | min_beginning = indices[hit_count] 429 | if min_beginning + 10 > end: 430 | max_ending = end 431 | else: 432 | max_ending = min_beginning + 10 433 | 434 | target = text[min_beginning:max_ending] 435 | 436 | for entry in input_types_list: 437 | if str(entry + "}") in target: 438 | used_input_types.append(entry) 439 | break 440 | 441 | hit_count += 1 442 | 443 | return used_input_types 444 | 445 | 446 | def get_tool_names(tool_list): 447 | """ SUMMARY: creates a list (set) of unique tool names for searching, auto-suggestion, etc. 448 | INPUT: a list of two-item tuple (tool, platform) 449 | OUTPUT: a de-duplicated list (set) of tool names """ 450 | 451 | tool_names = [] 452 | for x in tool_list: 453 | tool_names.append(x[0]) 454 | 455 | return set(tool_names) 456 | 457 | 458 | def initialize(settings): 459 | """ SUMMARY: installs PTF + toolkits, optionally downloads/loads user-supplied config file (overwriting existing) 460 | INPUT: three-item settings dictionary 461 | OUTPUT: no return values, modifies the filesystem and global config variable """ 462 | 463 | print("Initializing...") 464 | 465 | if isinstance(settings.get("tool"), str): 466 | config_yaml_url = str(settings.get("tool")) 467 | 468 | print("Downloading configuration file...") 469 | 470 | try: 471 | request = get(config_yaml_url) 472 | 473 | except RequestException: 474 | print_error(str("ERROR: Cannot download configuration file from " + config_yaml_url)) 475 | exit(1) 476 | 477 | else: 478 | if request.status_code == 200: 479 | try: 480 | remove(config.axiom.config_file) 481 | with open(config.axiom.config_file, 'wb') as config_file: 482 | config_file.write(request.content) 483 | 484 | except OSError: 485 | print_error("ERROR: Cannot replace existing configuration file") 486 | exit(1) 487 | 488 | else: 489 | config.axiom = config.AxiomConfig(config.axiom.config_file) 490 | setup_ptf() 491 | setup_toolkits() 492 | 493 | else: 494 | print_error("ERROR: Configuration file download failure") 495 | exit(1) 496 | 497 | elif settings.get("tool") is None: 498 | setup_ptf() 499 | setup_toolkits() 500 | 501 | else: 502 | print_error("ERROR: Invalid configuration file URL") 503 | exit(1) 504 | 505 | 506 | def load_commands(yam, inputs_pattern, input_types_list): 507 | """ SUMMARY: creates all command and action objects for a given tool file's YAML data 508 | INPUT: 1) a list of 2 dicts from the source YAML file, 2) a regex pattern (str), and 3) a list of strings 509 | OUTPUT: a two-item tuple of 1) a list of AxiomCommand objects and 2) a list of AxiomAction objects """ 510 | 511 | total = yam[1]['commands'].__len__() 512 | command_list = [] 513 | action_list = [] 514 | tool_string = "" 515 | 516 | i = 0 517 | while i < total: 518 | current_cmd = yam[1]['commands'][i] 519 | name = str(list(current_cmd.keys())[0]) 520 | 521 | for x in command_list: 522 | if x.name == name: 523 | tool_string = str(yam[0]["name"] + " (" + yam[0]["os"] + ") ") 524 | print_error(str("ERROR: " + tool_string + "contains non-unique command name \"" + name + "\"")) 525 | exit(1) 526 | 527 | for y in action_list: 528 | if y.name == name: 529 | tool_string = str(yam[0]["name"] + " (" + yam[0]["os"] + ") ") 530 | print_error(str("ERROR: " + tool_string + "contains non-unique action name \"" + name + "\"")) 531 | exit(1) 532 | 533 | prompt_type = str(list(list(current_cmd.values())[0][0].values())[0][0]) 534 | execution_type = str(list(list(current_cmd.values())[0][0].values())[0][1]) 535 | text = list(list(current_cmd.values())[0][1].values())[0] 536 | note = str(list(list(current_cmd.values())[0][4].values())[0]) 537 | 538 | raw_output_list = list(list(current_cmd.values())[0][3].values())[0] 539 | output_list = None 540 | if raw_output_list: 541 | output_list = load_outputs(raw_output_list, tool_string) 542 | 543 | raw_input_list = list(list(current_cmd.values())[0][2].values())[0] 544 | if raw_input_list: 545 | 546 | tokens, input_list = load_text_and_inputs(text, inputs_pattern, input_types_list, raw_input_list) 547 | command_list.append(AxiomCommand(name, prompt_type, execution_type, tokens, output_list, note, input_list)) 548 | 549 | else: 550 | action_list.append(AxiomAction(name, prompt_type, execution_type, text, output_list, note)) 551 | 552 | i += 1 553 | 554 | return command_list, action_list 555 | 556 | 557 | def load_inventory(): 558 | """ SUMMARY: instantiates the runtime toolkits that organize all tools and their commands/actions 559 | INPUT: none 560 | OUTPUT: a list of AxiomToolkit objects """ 561 | 562 | loadable_inventory_file = str(config.axiom.binary_folder + "/inventory.axiom") 563 | if path.exists(loadable_inventory_file): 564 | try: 565 | with open(loadable_inventory_file, 'rb') as inventory_dump: 566 | toolkits = load(inventory_dump) 567 | 568 | except (OSError, PickleError): 569 | print_error(str("ERROR: Failed to load inventory binary file " + loadable_inventory_file)) 570 | exit(1) 571 | else: 572 | return toolkits 573 | 574 | folders = [] 575 | if path.exists(config.axiom.inventory_folder): 576 | folders = listdir(config.axiom.inventory_folder) 577 | else: 578 | print_error(str("ERROR: Inventory folder " + config.axiom.inventory_folder + " not found")) 579 | exit(1) 580 | 581 | toolkits = [] 582 | 583 | for i in range(folders.__len__()): 584 | kit_name = folders[i] 585 | kit_folder = str(config.axiom.inventory_folder + "/" + kit_name) 586 | tool_list = [] 587 | 588 | for filename in listdir(kit_folder): 589 | current_file = str(kit_folder + "/" + filename) 590 | if current_file.endswith(".yml"): 591 | try: 592 | with open(current_file, 'r') as tool_file: 593 | tool_yaml = list(safe_load_all(tool_file))[0] 594 | tool_name = tool_yaml["name"] 595 | tool_platform = tool_yaml["os"] 596 | tool_list.append((tool_name, tool_platform)) 597 | 598 | except (OSError, parser.ParserError, scanner.ScannerError): 599 | print_error(str("ERROR: Failed to load " + current_file)) 600 | exit(1) 601 | 602 | tool_list = set(tool_list) 603 | toolkits.append(AxiomToolkit(kit_name, kit_folder, tool_list)) 604 | 605 | try: 606 | with open(loadable_inventory_file, 'wb') as inventory: 607 | dump(toolkits, inventory) 608 | 609 | except (OSError, PickleError): 610 | print_error(str("ERROR: Failed to save inventory binary file " + loadable_inventory_file)) 611 | exit(1) 612 | 613 | return toolkits 614 | 615 | 616 | def load_outputs(raw_output_list, tool): 617 | """ SUMMARY: retrieves a list of values representing each command/action output 618 | INPUT: 1) a list of outputs (str) taken directly from a YAML file 2) the target tool (str) 619 | OUTPUT: a list of two-item tuples """ 620 | 621 | output_list = [] 622 | 623 | output_count = 0 624 | 625 | try: 626 | while output_count < raw_output_list.__len__(): 627 | current_output = raw_output_list[output_count] 628 | 629 | if isinstance(current_output, dict): 630 | 631 | if list(current_output)[0] == "FILE": 632 | 633 | if list(list(current_output.values())[0].keys())[0] == "input": 634 | output_list.append(("F_INPUT", int(list(list(current_output.values())[0].values())[0]))) 635 | 636 | elif list(list(current_output.values())[0].keys())[0] == "string": 637 | output_list.append(("F_STRING", str(list(list(current_output.values())[0].values())[0]))) 638 | 639 | elif list(list(current_output.values())[0].keys())[0] == "prefix": 640 | input_number = int(list(list(current_output.values())[0].values())[0][0]) 641 | 642 | if isinstance(list(list(current_output.values())[0].values())[0][1], str): # single extension 643 | extension_string = str(list(list(current_output.values())[0].values())[0][1]) 644 | output_list.append(("F_PREFIX", (input_number, extension_string))) 645 | 646 | elif isinstance(list(list(current_output.values())[0].values())[0][1], list): # >1 extensions 647 | prefix_count = 0 648 | while prefix_count <= list(list(current_output.values())[0].values())[0][0]: 649 | extension_string = str(list(list( 650 | current_output.values())[0].values())[0][1][prefix_count]) 651 | output_list.append(("F_PREFIX", (input_number, extension_string))) 652 | prefix_count += 1 653 | 654 | elif list(current_output)[0] == "PROMPT": 655 | output_list.append(("PROMPT", str(list(current_output.values())[0]))) 656 | 657 | else: 658 | output_list.append(str(current_output)) 659 | 660 | output_count += 1 661 | 662 | except (AttributeError, IndexError, KeyError, TypeError, ValueError): 663 | print_error(str("ERROR: Invalid outputs defined for " + tool)) 664 | exit(1) 665 | 666 | return output_list 667 | 668 | 669 | def load_text_and_inputs(text, inputs_pattern, input_types_list, raw_input_list): 670 | """ SUMMARY: retrieves executable command/action text (tokens) and the inputs required at execution 671 | INPUT: 1) command text (str), 2) regex pattern (str), 3) list of types (list), and 4) list of inputs (list) 672 | OUTPUT: a two-item tuple containing 1) a list of strings and 2) a list of 2-item or 3-item tuples """ 673 | 674 | used_input_types = get_input_types(input_types_list, text) 675 | tokens = [] 676 | 677 | if isinstance(text, str): 678 | tokens = list(split(inputs_pattern, text)) 679 | elif isinstance(text, list): 680 | line_count = 0 681 | while line_count < text.__len__(): 682 | current_line = text[line_count] 683 | tokens.append(list(split(inputs_pattern, current_line))) 684 | line_count += 1 685 | 686 | input_list = [] 687 | input_count = 0 688 | while input_count < raw_input_list.__len__(): 689 | current_input = raw_input_list[input_count] 690 | current_type = used_input_types[input_count] 691 | if isinstance(current_input, str): 692 | input_list.append(tuple((current_input, current_type))) 693 | elif isinstance(current_input, dict): 694 | current_name = list(current_input.keys())[0] 695 | current_options = list(current_input.values())[0] 696 | input_list.append(tuple((current_name, current_type, current_options))) 697 | 698 | input_count += 1 699 | 700 | return tokens, input_list 701 | 702 | 703 | def load_tool_list(inventory): 704 | """ SUMMARY: creates a de-duplicated list of all tools present in all toolkits 705 | INPUT: a list of AxiomToolkit objects 706 | OUTPUT: a list of two-item tuples (tool, platform) """ 707 | 708 | loadable_list_file = str(config.axiom.binary_folder + '/tool_list.axiom') 709 | if path.exists(loadable_list_file): 710 | try: 711 | with open(loadable_list_file, 'rb') as list_dump: 712 | loaded_list = load(list_dump) 713 | 714 | except (OSError, PickleError): 715 | print_error(str("ERROR: Failed to load tool list binary file " + loadable_list_file)) 716 | exit(1) 717 | 718 | else: 719 | return loaded_list 720 | 721 | master_tool_list = [] 722 | 723 | for i in range(inventory.__len__()): 724 | for x in inventory[i].tool_name_list: 725 | master_tool_list.append(x) 726 | 727 | master_tool_list = set(master_tool_list) 728 | 729 | try: 730 | with open(loadable_list_file, 'wb') as tool_list: 731 | dump(list(master_tool_list), tool_list) 732 | 733 | except (OSError, PickleError): 734 | print_error(str("ERROR: Failed to save tool list binary file " + loadable_list_file)) 735 | exit(1) 736 | 737 | return list(master_tool_list) 738 | 739 | 740 | def load_tools(inventory, unloaded_tools): 741 | """ SUMMARY: imports all tool data from all YAML files from all inventory folders 742 | INPUT: 1) a list of AxiomToolkit objects, and 2) a list of two-item tuples (tool, platform) 743 | OUTPUT: a list of AxiomTool objects """ 744 | 745 | loadable_tools_file = str(config.axiom.binary_folder + "/tools.axiom") 746 | if path.exists(loadable_tools_file): 747 | try: 748 | with open(loadable_tools_file, 'rb') as tools_dump: 749 | loaded_tools = load(tools_dump) 750 | 751 | except (OSError, PickleError): 752 | print_error(str("ERROR: Failed to load tools binary file " + loadable_tools_file)) 753 | exit(1) 754 | 755 | else: 756 | return loaded_tools 757 | 758 | tools = [] 759 | 760 | for i in range(len(inventory)): 761 | folder = inventory[i].location 762 | 763 | for filename in listdir(folder): 764 | current_file = str(folder + "/" + filename) 765 | 766 | if current_file.endswith(".yml"): 767 | try: 768 | with open(current_file, 'r') as tool_file: 769 | tool = list(safe_load_all(tool_file)) 770 | current_tool = (tool[0]["name"], tool[0]["os"]) 771 | 772 | if current_tool in unloaded_tools: 773 | command_list, action_list = load_commands(tool, config.axiom.inputs_pattern, 774 | config.axiom.input_types_list) 775 | tools.append(AxiomTool(tool[0]["name"], tool[0]["os"], tool[0]["ptf_module"], 776 | tool[0]["description"], action_list, command_list)) 777 | unloaded_tools.remove(current_tool) 778 | else: 779 | tool_id = resolve_tool_id(current_tool, tools) 780 | if able_to_merge(tool, tool_id, tools): 781 | if merge(tool, tool_id, tools, config.axiom.inputs_pattern, 782 | config.axiom.input_types_list): 783 | continue 784 | else: 785 | print_error(str("ERROR: Merge failure for " + str(tool[0]["name"]) + " from " + 786 | str(current_file))) 787 | exit(1) 788 | else: 789 | print_error(str("ERROR: Unable to merge " + str(tool[0]["name"]) + " from " + 790 | str(current_file))) 791 | exit(1) 792 | 793 | except (AttributeError, IndexError, KeyError, OSError, TypeError, ValueError): 794 | print_error(str("ERROR: Failed to load " + current_file)) 795 | exit(1) 796 | 797 | for item in tools: 798 | item.initialize_combined_list() 799 | 800 | try: 801 | with open(loadable_tools_file, 'wb') as axiom: 802 | dump(tools, axiom) 803 | 804 | except (OSError, PickleError): 805 | print_error(str("ERROR: Failed to save tools binary file " + loadable_tools_file)) 806 | exit(1) 807 | 808 | return tools 809 | 810 | 811 | def merge(tool, tool_id, tools, inputs_pattern, input_types_list): 812 | """ SUMMARY: merges new commands/actions into existing AxiomTool objects 813 | INPUT: 1) list of two dictionaries 2) tool ID value (int) 3) list of AxiomTool objects 814 | 4) regex pattern (str) and 5) list of strings 815 | OUTPUT: Returns True after completing merge procedure """ 816 | 817 | command_list, action_list = load_commands(tool, inputs_pattern, input_types_list) 818 | action_count = 0 819 | command_count = 0 820 | 821 | if action_list.__len__() > 0: 822 | while action_count < action_list.__len__(): 823 | tools[tool_id].action_list.append(action_list[action_count]) 824 | action_count += 1 825 | 826 | if command_list.__len__() > 0: 827 | while command_count < command_list.__len__(): 828 | tools[tool_id].command_list.append(command_list[command_count]) 829 | command_count += 1 830 | 831 | return True 832 | 833 | 834 | def new_generate_command(): 835 | """ SUMMARY: prompts user with data entry questions and prints a complete and valid YAML snippet to the screen 836 | INPUT: none 837 | OUTPUT: none, prints to the screen and exits """ 838 | 839 | name = prompt("[AXIOM] Enter command name: ") 840 | 841 | prompt_selection = new_get_prompt_selection() 842 | execution_type = new_get_execution_type(prompt_selection) 843 | text = new_get_text() 844 | inputs = new_get_inputs(text) 845 | outputs = new_get_outputs(execution_type, text) 846 | 847 | note = prompt("[AXIOM] Enter command note: ") 848 | 849 | name = new_get_escaped_text(name) 850 | note = new_get_escaped_text(note) 851 | 852 | new_print_finalized_command_text(name, prompt_selection, execution_type, text, inputs, outputs, note) 853 | 854 | exit(0) 855 | 856 | 857 | def new_get_escaped_text(text): 858 | """ SUMMARY: replaces any backslash and double-quote characters with backslash-escaped character sequences 859 | INPUT: command text line(s) (list or str) 860 | OUTPUT: returns backslash-escaped command text (list or str) """ 861 | 862 | if isinstance(text, list): 863 | new_list = [] 864 | for line in range(text.__len__()): 865 | new_list.append(text[line].replace("\\", "\\\\").replace("\"", "\\\"")) 866 | return new_list 867 | 868 | else: 869 | return text.replace("\\", "\\\\").replace("\"", "\\\"") 870 | 871 | 872 | def new_get_execution_type(prompt_selection): 873 | """ SUMMARY: prompts user to enter the command execution type 874 | INPUT: the current command's prompt type (str) 875 | OUTPUT: returns the execution type name (str) """ 876 | 877 | if prompt_selection == "other": 878 | return "NX" 879 | 880 | print("\nExecution Types\n") 881 | 882 | print(" 1\tstandalone") 883 | print(" 2\tautonomous") 884 | print(" 3\tinteractive") 885 | print(" 4\tNX") 886 | 887 | number = prompt("\n[AXIOM] Select an option: ") 888 | 889 | try: 890 | number = int(number) 891 | if number == 1: 892 | return "standalone" 893 | elif number == 2: 894 | return "autonomous" 895 | elif number == 3: 896 | return "interactive" 897 | elif number == 4: 898 | return "NX" 899 | else: 900 | print_error("ERROR: Invalid execution type selection") 901 | exit(1) 902 | 903 | except (ValueError, TypeError): 904 | print_error("ERROR: Invalid execution type selection") 905 | exit(1) 906 | 907 | 908 | def new_get_inputs(text): 909 | """ SUMMARY: prompts user to enter input descriptions and related data 910 | INPUT: the command text (list or str) 911 | OUTPUT: returns the inputs text line (str) """ 912 | 913 | inputs = "[" 914 | 915 | used_input_types = get_input_types(config.axiom.input_types_list, text) 916 | input_count = used_input_types.__len__() 917 | 918 | for i in range(input_count): 919 | description = prompt(str("[AXIOM] Enter name for input " + 920 | str("(" + str(i + 1) + "/" + str(input_count) + ")") + 921 | " {" + used_input_types[i] + "}: ")) 922 | description = new_get_escaped_text(description) 923 | 924 | if used_input_types[i] in ["INTMENU", "STRMENU"]: 925 | option_count = prompt(str("[AXIOM] Enter number of \"" + description + "\" options: ")) 926 | try: 927 | option_count = int(option_count) 928 | if option_count <= 0: 929 | print_error("ERROR: Invalid number of options") 930 | exit(1) 931 | option_text = "[" 932 | except (ValueError, TypeError): 933 | print_error("ERROR: Invalid number of options") 934 | exit(1) 935 | 936 | if used_input_types[i] == "INTMENU": 937 | for x in range(option_count): 938 | single_option = prompt(str("[AXIOM] Enter \"" + description + "\" option (" + 939 | str(x + 1) + "/" + str(option_count) + ") {INT}: ")) 940 | try: 941 | single_option = int(single_option) 942 | except (ValueError, TypeError): 943 | print_error("ERROR: Invalid integer option") 944 | exit(1) 945 | 946 | option_text = str(option_text + str(single_option) + ",") 947 | option_text = str(option_text[:-1] + "]") 948 | 949 | elif used_input_types[i] == "STRMENU": 950 | for x in range(option_count): 951 | single_option = prompt(str("[AXIOM] Enter \"" + description + "\" option (" + 952 | str(x + 1) + "/" + str(option_count) + ") {STR}: ")) 953 | single_option = new_get_escaped_text(single_option) 954 | 955 | option_text = str(option_text + "\"" + str(single_option) + "\"" + ",") 956 | option_text = str(option_text[:-1] + "]") 957 | 958 | inputs = str(inputs + "{\"" + description + "\":" + option_text + "},") 959 | 960 | else: 961 | 962 | inputs = str(inputs + "\"" + description + "\",") 963 | 964 | if inputs == "[": 965 | return "null" 966 | else: 967 | inputs = str(inputs[:-1] + "]") 968 | return inputs 969 | 970 | 971 | def new_get_output_details(input_count, current_output_index, output_count): 972 | """ SUMMARY: prompts user to select output type and enter type-specific details 973 | INPUT: 1) number of total inputs (int) 2) current output number (int) 3) total number of outputs (int) 974 | OUTPUT: returns the outputs text line (str) """ 975 | 976 | print(str("[AXIOM] Select output type for remaining output (" + 977 | str(current_output_index + 1) + "/" + str(output_count) + "): ")) 978 | print("\nOutput Types\n") 979 | 980 | print(" 1\tFile (input)\tfilename is entirely user-controlled command input") 981 | print(" 2\tFile (prefix)\tfilename prefix is command input, file extension(s) hardcoded") 982 | print(" 3\tFile (string)\tfilename is entirely hardcoded") 983 | print(" 4\tSTDERR\t\tstandard error") 984 | 985 | number = prompt("\n[AXIOM] Select an option: ") 986 | 987 | try: 988 | number = int(number) 989 | if number == 1: 990 | if input_count <= 0: 991 | print_error("ERROR: Output type requires at least one command input") 992 | exit(1) 993 | input_number = prompt("[AXIOM] Enter the corresponding input number: ") 994 | input_number = int(input_number) 995 | if int(input_number - 1) in range(input_count): 996 | return str("{\"FILE\":{\"input\":" + str(input_number) + "}}") 997 | else: 998 | print_error("ERROR: Invalid input number") 999 | exit(1) 1000 | 1001 | elif number == 2: 1002 | if input_count <= 0: 1003 | print_error("ERROR: Output type requires at least one command input") 1004 | exit(1) 1005 | extensions = "" 1006 | entry = prompt("[AXIOM] Enter the corresponding input number: ") 1007 | entry = int(entry) 1008 | if int(entry - 1) in range(input_count): 1009 | extension_count = prompt("[AXIOM] Enter number of file extensions: ") 1010 | extension_count = int(extension_count) 1011 | if extension_count > 0: 1012 | for e in range(extension_count): 1013 | current_ext = prompt("[AXIOM] Enter file extension (" + 1014 | str(e + 1) + "/" + str(extension_count) + "): ") 1015 | current_ext = new_get_escaped_text(current_ext) 1016 | extensions = str(extensions + "\"" + current_ext + "\",") 1017 | extensions = extensions[:-1] 1018 | return str("{\"FILE\":{\"prefix\":[" + str(entry) + "," + extensions + "]}}") 1019 | else: 1020 | print_error("ERROR: Invalid number of file extensions") 1021 | exit(1) 1022 | else: 1023 | print_error("ERROR: Invalid input number") 1024 | exit(1) 1025 | 1026 | elif number == 3: 1027 | filename = prompt("[AXIOM] Enter the output filename: ") 1028 | filename = new_get_escaped_text(filename) 1029 | return str("{\"FILE\":{\"string\":\"" + str(filename) + "\"}}") 1030 | 1031 | elif number == 4: 1032 | return "\"STDERR\"" 1033 | 1034 | else: 1035 | print_error("ERROR: Invalid output type selection") 1036 | exit(1) 1037 | 1038 | except (ValueError, TypeError): 1039 | print_error("ERROR: Invalid number entered") 1040 | exit(1) 1041 | 1042 | 1043 | def new_get_outputs(execution_type, text): 1044 | """ SUMMARY: prompts user to enter output data 1045 | INPUT: 1) command execution type name (str) 2) command text (list or str) 1046 | OUTPUT: returns completed outputs text line (str) """ 1047 | 1048 | input_count = get_input_types(config.axiom.input_types_list, text).__len__() 1049 | outputs = "[" 1050 | 1051 | answer = prompt("[AXIOM] Does command output to STDOUT? [Y/n] ") 1052 | if answer not in ["Y", "y", "Yes", "yes"]: 1053 | pass 1054 | else: 1055 | outputs = "[\"STDOUT\"," 1056 | 1057 | if execution_type == "interactive": 1058 | print("[AXIOM] Select prompt type emitted by interactive command: ") 1059 | prompt_type = new_get_prompt_selection() 1060 | if outputs == "[\"STDOUT\",": 1061 | outputs = str("[\"STDOUT\"," + "{\"PROMPT\":\"" + prompt_type + "\"},") 1062 | else: 1063 | outputs = str("[{\"PROMPT\":\"" + prompt_type + "\"},") 1064 | 1065 | output_count = prompt("[AXIOM] Enter number of remaining outputs: ") 1066 | 1067 | try: 1068 | output_count = int(output_count) 1069 | 1070 | except (ValueError, TypeError): 1071 | print_error("ERROR: Invalid number of outputs") 1072 | exit(1) 1073 | 1074 | if output_count <= 0: 1075 | if outputs == "[": 1076 | return "null" 1077 | else: 1078 | return str(outputs[:-1] + "]") 1079 | 1080 | for y in range(output_count): 1081 | outputs = str(outputs + new_get_output_details(input_count, y, output_count) + ",") 1082 | 1083 | return str(outputs[:-1] + "]") 1084 | 1085 | 1086 | def new_get_prompt_selection(): 1087 | """ SUMMARY: prompts user to select command prompt type 1088 | INPUT: none, gets input from user 1089 | OUTPUT: returns a prompt name from the global config (str) """ 1090 | 1091 | print("\nPrompts\n") 1092 | 1093 | for i in range(config.axiom.prompts.__len__()): 1094 | print(" " + str(i + 1) + "\t" + str(config.axiom.prompts[i][0])) 1095 | 1096 | number = prompt("\n[AXIOM] Select an option: ") 1097 | 1098 | try: 1099 | number = int(number) 1100 | if int(number - 1) in range(config.axiom.prompts.__len__()): 1101 | return config.axiom.prompts[number - 1][0] 1102 | 1103 | else: 1104 | print_error("ERROR: Invalid prompt selection") 1105 | exit(1) 1106 | 1107 | except (ValueError, TypeError): 1108 | print_error("ERROR: Invalid prompt selection") 1109 | exit(1) 1110 | 1111 | 1112 | def new_get_text(): 1113 | """ SUMMARY: prompts user for number of command text lines and the line contents 1114 | INPUT: none, gets input from the user 1115 | OUTPUT: returns completed and sanitized command text (list or str) """ 1116 | 1117 | line_count = prompt("[AXIOM] Enter number of text input lines: ") 1118 | 1119 | try: 1120 | line_count = int(line_count) 1121 | except (ValueError, TypeError): 1122 | print_error("ERROR: Invalid number of lines") 1123 | exit(1) 1124 | 1125 | if line_count <= 0: 1126 | print_error("ERROR: Invalid number of lines") 1127 | exit(1) 1128 | 1129 | elif line_count == 1: 1130 | text = prompt("[AXIOM] Enter command text: ") 1131 | 1132 | else: 1133 | text = [] 1134 | for i in range(line_count): 1135 | text.append(prompt(str("[AXIOM] Enter command text (line " + str(i + 1) + "): "))) 1136 | 1137 | return new_get_escaped_text(text) 1138 | 1139 | 1140 | def new_print_finalized_command_text(name, prompt_selection, execution_type, text, inputs, outputs, note): 1141 | """ SUMMARY: prints newly-generated YAML text to the screen 1142 | INPUT: seven variables generated by related functions, all are strings but "text" can also be a list 1143 | OUTPUT: none, only prints to the screen """ 1144 | 1145 | print() 1146 | 1147 | print(" - \"" + name + "\":") 1148 | print(" - type: [\"" + prompt_selection + "\",\"" + execution_type + "\"]") 1149 | 1150 | if isinstance(text, str): 1151 | print(" - text: \"" + text + "\"") 1152 | else: 1153 | print(" - text:") 1154 | for i in range(text.__len__()): 1155 | print(" - \"" + text[i] + "\"") 1156 | 1157 | print(" - input: " + inputs + "") 1158 | print(" - output: " + outputs + "") 1159 | print(" - note: \"" + note + "\"") 1160 | 1161 | print() 1162 | 1163 | 1164 | def print_banner(banner_file): 1165 | """ SUMMARY: displays ASCII art from file and other introductory info 1166 | INPUT: filename (str) of text file on filesystem 1167 | OUTPUT: none, only prints to screen """ 1168 | 1169 | try: 1170 | with open(banner_file, 'r') as file: 1171 | data = file.readlines() 1172 | for line in data: 1173 | print(Fore.RED + line.replace('\n', '')) 1174 | 1175 | print(Style.RESET_ALL, end='') 1176 | 1177 | except OSError: 1178 | print_error(str("ERROR: Unable to access banner file " + banner_file)) 1179 | exit(1) 1180 | 1181 | else: 1182 | print(" C9EE FD5E 15DA 9C02 1B0C 603C A397 0118 D56B 2E35 ") 1183 | print() 1184 | print(" Created by Mike Iacovacci https://payl0ad.run ") 1185 | print() 1186 | print() 1187 | 1188 | 1189 | def print_stats(inventory, tool_list, tools): 1190 | """ SUMMARY: displays counts of loaded tools, commands/actions, and toolkits 1191 | INPUT: 1) list of AxiomToolkit objects objects 2) de-deplicated list of tuples 3) list of AxiomTool objects 1192 | OUTPUT: none, only prints to the screen """ 1193 | 1194 | action_count = 0 1195 | command_count = 0 1196 | for tool in tools: 1197 | current_actions = tool.action_list.__len__() 1198 | current_commands = tool.command_list.__len__() 1199 | action_count = action_count + current_actions 1200 | command_count = command_count + current_commands 1201 | 1202 | combined_count = str(action_count + command_count) 1203 | 1204 | tool_count = str(tool_list.__len__()) 1205 | toolkit_count = str(inventory.__len__()) 1206 | 1207 | print("\n" + "Loaded " + 1208 | combined_count + " commands for " + 1209 | tool_count + " unique tools from " + 1210 | toolkit_count + " toolkits." 1211 | "\n") 1212 | 1213 | 1214 | def reload(): 1215 | """ SUMMARY: deletes and recreates binary folder causing all YAML tool files to be deserialized again 1216 | INPUT: none 1217 | OUTPUT: none """ 1218 | 1219 | print("Reloading...") 1220 | 1221 | delete_and_recreate_folder(config.axiom.binary_folder) 1222 | 1223 | 1224 | def resolve_tool_id(potential_tool, tools): 1225 | """ SUMMARY: searches for a tool's ID number using a user-supplied tool name string 1226 | INPUT: 1) a tool name (str), and 2) a list of AxiomTool objects 1227 | OUTPUT: a tool ID value (int) or -1 if no match is found """ 1228 | 1229 | tool_id = 0 1230 | while tool_id < tools.__len__(): 1231 | if tools[tool_id].name == potential_tool[0] and tools[tool_id].platform == potential_tool[1]: 1232 | return tool_id 1233 | tool_id += 1 1234 | 1235 | return -1 1236 | 1237 | 1238 | def set_user_expectations(settings): 1239 | """ SUMMARY: prints a message so the user expects to wait while the YAML is deserialized 1240 | INPUT: three-item settings dictionary 1241 | OUTPUT: no return value, only prints to the screen conditionally """ 1242 | 1243 | if path.exists(str(config.axiom.binary_folder + "/inventory.axiom")) and \ 1244 | path.exists(str(config.axiom.binary_folder + "/tool_list.axiom")) and \ 1245 | path.exists(str(config.axiom.binary_folder + "/tools.axiom")) or \ 1246 | settings.get("mode") in ["init", "reload"]: 1247 | return 1248 | else: 1249 | print("Initializing...") 1250 | 1251 | 1252 | def setup_folders(settings): 1253 | """ SUMMARY: initializes folders for history/binary files and installs PTF if missing 1254 | INPUT: three-item settings dictionary 1255 | OUTPUT: none, causes filesystem modifications within user-defined locations """ 1256 | 1257 | set_user_expectations(settings) 1258 | 1259 | create_missing_folder(config.axiom.history_folder) 1260 | create_missing_folder(config.axiom.binary_folder) 1261 | 1262 | if not path.exists(config.axiom.ptf_folder): 1263 | setup_ptf() 1264 | if not path.exists(config.axiom.inventory_folder): 1265 | setup_toolkits() 1266 | 1267 | 1268 | def setup_ptf(): 1269 | """ SUMMARY: deletes existing PTF folder and downloads/installs the latest version from GitHub master branch 1270 | INPUT: none 1271 | OUTPUT: no return values, modifies the filesystem """ 1272 | 1273 | download_and_extract_zip("https://github.com/trustedsec/ptf/archive/master.zip", 1274 | "ptf-master", 1275 | config.axiom.ptf_folder, 1276 | "The PenTesters Framework (PTF)") 1277 | 1278 | 1279 | def setup_toolkits(): 1280 | """ SUMMARY: deletes existing inventory folder, downloads all listed toolkits, and reloads the binary data 1281 | INPUT: none 1282 | OUTPUT: no return values, modifies the filesystem """ 1283 | 1284 | delete_and_recreate_folder(config.axiom.inventory_folder) 1285 | 1286 | for toolkit in config.axiom.toolkits: 1287 | download_and_extract_zip(toolkit[2], 1288 | toolkit[1], 1289 | str(config.axiom.inventory_folder + "/" + toolkit[0]), 1290 | toolkit[0]) 1291 | 1292 | reload() 1293 | 1294 | 1295 | def tool_selection_prompt(tool_list, tool_names, tools): 1296 | """ SUMMARY: prompts user to select a tool, provides a fuzzy word completer interface 1297 | INPUT: 1) list of two-item tuples (name, platform), 2) set of tool names, and 3) list of AxiomTool objects 1298 | OUTPUT: exit value (int) """ 1299 | 1300 | tool_names = FuzzyCompleter(WordCompleter(tool_names)) 1301 | 1302 | completer_style = ptkStyle.from_dict({ 1303 | "completion-menu": "bg:#111111", 1304 | "scrollbar.background": "bg:#111111", 1305 | "scrollbar.button": "bg:#999999", 1306 | "completion-menu.completion.current": "nobold bg:ansired", 1307 | "completion-menu.completion fuzzymatch.outside": "nobold fg:#AAAAAA", 1308 | "completion-menu.completion fuzzymatch.inside": "nobold fg:ansired", 1309 | "completion-menu.completion fuzzymatch.inside.character": "nobold nounderline fg:ansired", 1310 | "completion-menu.completion.current fuzzymatch.outside": "nobold fg:#AAAAAA", 1311 | "completion-menu.completion.current fuzzymatch.inside": "nobold fg:#AAAAAA", 1312 | "completion-menu.completion.current fuzzymatch.inside.character": "nobold nounderline fg:#AAAAAA"}) 1313 | 1314 | while True: 1315 | text = prompt('[AXIOM] Enter tool: ', completer=tool_names, complete_while_typing=True, style=completer_style) 1316 | 1317 | if text == "exit" or text == "quit": 1318 | return 0 1319 | if text == "": 1320 | continue 1321 | 1322 | tool_id = disambiguate_tool_name(text, tool_list, tools) 1323 | if tool_id < 0: 1324 | print_error("ERROR: Invalid tool name") 1325 | else: 1326 | tool = tools[tool_id] 1327 | command_selection_prompt(tool) 1328 | 1329 | return 1 1330 | 1331 | 1332 | def validate_privileges(mode): 1333 | """ SUMMARY: confirms effective root privilege level if writing to the filesystem or spawning a subprocess 1334 | INPUT: program mode type (str) 1335 | OUTPUT: none """ 1336 | 1337 | if mode not in ["show", "new"]: 1338 | if geteuid() != 0: 1339 | print_error("ERROR: AXIOM requires root privileges") 1340 | exit(1) 1341 | --------------------------------------------------------------------------------