converting a python script into object oriented program

10.3k Views Asked by At

Hello I am a semi self-taught student, looking to increase my skills, made couples of script on Python using different APIs. I would like to at some point use this script as a module in a bigger project of mine. I have as idea to turn this script in a more object oriented program type of programming. with classes and functions to maybe call from an other file. I started creating the class and function but don't think I am doing it the right way I am now at the part with all my prints and dont know how to convert a script into OOP the right way, any tips or anything that would guide me a little bit in doing it the right way would be appreciated, the first lines are what I started converting into OOP.

The script works fine, it basically asked a question and depending on the input from the user it build the url to make a request to Vulners website and then loop over the nested dictionnaries and print out the keys and values i was interested in seeing, it also save the data locally into a json file, not really usefull for the moment. any tips well apreciated!

class vulnersApi(object):
    def type_of_search(self, software, bulletin, inp):
        self.bulletin = software
        self.collection = bulletin
        self.inp = inp

        while inp != "software" and inp!= "bulletin":
            inp = input("You must choose your type of research(software or bulletin: ")
            if inp != "software" and inp != "bulletin":
                return "You must choose your type of research"

    def search_software(self, myqueryName, sortedHow, numberResults, numberSkipped):
        self.myqueryName = myqueryName
        self.sortedHow = sortedHow
        self.numberResults = numberResults
        self.numberSkipped = numberSkipped

        if self.inp == "software":
            myqueryName = input("What would you like to search? ")
            sortedHow = input("How should we sort the results? (published, cvss.score) ")
            sortedHow = input("How should we sort the results? (published, cvss.score) ")
            numberSkipped = input("How many results do you want to skip? ")
            new_url = "{}?query=affectedSoftware.name%3A{}&sort={}&size={}&skip={}".format(URL, myqueryName, sortedHow, numberResults, numberSkipped)

    def search_bulletin(self, typeSearch, sortedHow, numberResults, numberSkipped):
        self.typeSearch = typeSearch
        self.sortedHow = sortedHow
        self.numberResults = numberResults
        self.numberSkipped = numberSkipped

        if self.inp == "bulletin":
            typeSearch = input("Which db you want the info from? (cve, exploitdb, osvdb, openvas, securityvulns, nessus, metasploit, centos, malwarebytes, symantec, etc...) ")
            sortedHow = input("How should we sort the results? (published, cvss.score) ")
            numberResults = input("How many results? ")
            numberSkipped = input("How many results do you want to skip? ")
            new_url = "{}?query=type%3A{}&sort={}&size={}&skip={}".format(URL, typeSearch, sortedHow, numberResults, numberSkipped)

    def url_request(self):
        response = requests.get(new_url)
        response.raise_for_status()
        json_string = json.dumps(response.text)
        mydata = json.loads(response.text)

Script is bellow

# base url

URL = "https://vulners.com/api/v3/search/lucene/"

# choose between 2 types of research
inp = ""
while inp != "software" and inp != "bulletin" and inp !="collection":
    inp = input("You must chose your type of research(software or bulletin): ")
    if inp != "software" and inp != "bulletin" and inp !="collection":
        print("you must chose your type of research")

print("-"*30)
# if chosed software heres the questions asked to build the full url
if inp == "software":
    myqueryName = input("What would you like to search? ")
    sortedHow = input("How should we sort the results? (published, cvss.score) ")
    numberResults = input("How many results? ")
    numberSkipped = input("How many results do you want to skip? ")
    new_url = "{}?query=affectedSoftware.name%3A{}&sort={}&size={}&skip={}".format(URL, myqueryName, sortedHow, numberResults, numberSkipped)
# if chosed bulletin heres the questions asked to build the full url
if inp == "bulletin":
    typeSearch = input("Which db you want the info from? (cve, exploitdb, osvdb, openvas, securityvulns, nessus, metasploit, centos, malwarebytes, symantec, etc...) ")
    sortedHow = input("How should we sort the results? (published, cvss.score) ")
    numberResults = input("How many results? ")
    numberSkipped = input("How many results do you want to skip? ")
    new_url = "{}?query=type%3A{}&sort={}&size={}&skip={}".format(URL, typeSearch, sortedHow, numberResults, numberSkipped)

# making the request and converting the json
resp = requests.get(new_url)
resp.raise_for_status()
json_string=json.dumps(resp.text)
mydata=json.loads(resp.text)
# to see the url
print(resp.url)
print("-"*30)


# loop to go through nested dictionnary returned by the requests
mydata2 = mydata['data']
for mydata3 in mydata2['search']:
#     mydata4 = mydata3['_source']
#     mydata5 = mydata3['highlight']
    for mydata4 in mydata['data']:
        print("Title: {}".format(mydata3['_source']['title']))
        print("Index: {}".format(mydata3['_index']))
        print("Type: {}".format(mydata3['_type']))
        print("Id: {}".format(mydata3['_id']))
        print("Type: {}".format(mydata3['_source']['type']))
        print("Cvss: {}".format(mydata3['_source']['cvss']))
        print("Flat description: {}".format(mydata3['flatDescription']))
        print("Bulletin family: {}".format(mydata3['_source']['bulletinFamily']))
        print("Description: {}".format(mydata3['_source']['description']))
        print("Vhref: {}".format(mydata3['_source']['vhref']))
        print("Href: {}".format(mydata3['_source']['href']))
        print("Id: {}".format(mydata3['_source']['id']))
        print("Lastseen: {}".format(mydata3['_source']['lastseen']))
        print("Modified: {}".format(mydata3['_source']['modified']))
        print("Published: {}".format(mydata3['_source']['published']))
        print("-"*30)
    # saving the data locally inside a json file
        with open('testfile2.json', 'w') as fd:
            print(json.dump(resp.text, fd, indent=2))
            fd.write(resp.text)
1

There are 1 best solutions below

1
On

I would like to at some point use this script as a module in a bigger project of mine. I have as idea to turn this script in a more object oriented program type of programming. with classes and functions to maybe call from an other file.

There are some things I want to note first:

  • You don't need your code to be object-oriented in order to reuse it somewhere else, just put all the useful functions in one python file and import it somewhere else.
  • It only makes sense to refactor code that can make use of the advantages of OOP (like multiple instances with different values but shared functions)

I had a go at refactoring you code:

import requests
import json

class VulnersResult:
    def __init__(self, json_data):
        self.title = json_data['_source']['title']
        self.type = json_data['_type']
        self.id = json_data['_id']
        self.source_type = json_data['_source']['type']
        self.cvss = json_data['_source']['cvss']
        self.flat_description = json_data['flatDescription']
        self.bulletin_family = json_data['_source']['bulletinFamily']
        self.description = json_data['_source']['description']
        self.vhref = json_data['_source']['vhref']
        self.href = json_data['_source']['href']
        self.source_id = json_data['_source']['id']
        self.lastseen = json_data['_source']['lastseen']
        self.modified = json_data['_source']['modified']
        self.published = json_data['_source']['published']
    def __str__(self):
        lines = ["Title: {}".format(self.title),
                 "Type: {}".format(self.type),
                 "Id: {}".format(self.id),
                 "Type: {}".format(self.source_type),
                 "Cvss: {}".format(self.cvss),
                 "Flat description: {}".format(self.flat_description),
                 "Bulletin family: {}".format(self.bulletin_family),
                 "Description: {}".format(self.description),
                 "Vhref: {}".format(self.vhref),
                 "Href: {}".format(self.href),
                 "Id: {}".format(self.source_id),
                 "Lastseen: {}".format(self.lastseen),
                 "Modified: {}".format(self.modified),
                 "Published: {}".format(self.published)]
        return "\n".join(lines)

class VulnersAPI:
    BASE_URL = "https://vulners.com/api/v3/search/lucene/"
    def __init__(self, research_type, query, sorted_by, results_count, skip_count):
        if research_type == "software":
            self.query_url = "{}?query=affectedSoftware.name%3A{}&sort={}&size={}&skip={}".format(self.BASE_URL, query, sorted_by, results_count, skip_count)
        elif research_type == "bulletin":
            self.query_url = "{}?query=type%3A{}&sort={}&size={}&skip={}".format(self.BASE_URL, query, sorted_by, results_count, skip_count)
        else:
            raise RuntimeError("{} is not a valid research type. research_type must be 'software' or 'bulletin'".format(research_type))
        response = requests.get(self.query_url)
        response.raise_for_status()
        self.raw_data = response.json()
        self.results = [VulnersResult(data) for data in self.raw_data['data']['search']]
    def result_text(self):
        return ("\n"+("-"*30)+"\n").join([str(result) for result in self.results])

if __name__ == "__main__":
    # choose between 2 types of research
    inp = ""
    while inp != "software" and inp != "bulletin" and inp !="collection":
        inp = input("You must chose your type of research(software or bulletin): ")
        if inp != "software" and inp != "bulletin" and inp !="collection":
            print("you must chose your type of research")

    print("-"*30)
    if inp == "software":
        # if "software" was chosen, ask these additional questions
        query = input("What would you like to search? ")
        sorted_by = input("How should we sort the results? (published, cvss.score) ")
        results_count = input("How many results? ")
        skip_count = input("How many results do you want to skip? ")
        api = VulnersAPI("software", query, sorted_by, results_count, skip_count)
    if inp == "bulletin":
        # if "bulletin" was chosen, ask these additional questions
        query = input("Which db you want the info from? (cve, exploitdb, osvdb, openvas, securityvulns, nessus, metasploit, centos, malwarebytes, symantec, etc...) ")
        sorted_by = input("How should we sort the results? (published, cvss.score) ")
        results_count = input("How many results? ")
        skip_count = input("How many results do you want to skip? ")
        api = VulnersAPI("software", query, sorted_by, results_count, skip_count)
    print()

    with open('testfile2.json', 'w') as fd:
        json.dump(api.raw_data, fd)

    print(api.result_text())

The first step was to abstract away any code that handles user input to make the library compatible with all kinds of usecases. The second step was to think about the different entities that deserve their own class. I chose VulnersResult and VulnersAPI, because I think it makes sense to have them as separate entities and because you are able to add interesting functions later. (Something like to_table_row() or is_affecting_my_system() for VulnersResult and filter_by_type() or sort_by_lastseen() for VulnersAPI). In a third step you have to decide what properties each class should have. Properties should be variables that are used in multiple functions or variables that the user of a class should have access to. For example I did not add search_type or query as properties of VulnersAPI, because they are only used to generate the url, but if there is someone that uses these values, they should be stored as properties.

To use your python script as you are currently doing and to provide example code to the users of the library, I added the original code that handles the user input after if __name__ == "__main__". When you start the script directly it will ask the user for input, just as it does now, but when you import it, the script ignores the last part and only imports the definitions of VulnerResult and VulnerAPI.

I hope this example together with the steps to get there are helpful to you and enable you to tackle your big problem :)

Lastly, here are some resources that I think every python programmer should have read

and if you are programming with the latest version of python you can use f-strings to make string formatting even better.