4
\$\begingroup\$

I'm starting writing in Python and is tasked with writing a simple tool that opens a web-page, collects all the URLs (both complete like http://mail.com and relative ones like /mail/inbox) and returns a list of them prefixed with the domain name. The code is below.

Could someone take a look and advise on what could've been better in terms of the code style, readability, general coding principles and code quality?

Appreciate your help!

import sys, requests
from bs4 import BeautifulSoup
from urllib.parse import urlparse

def main(argv):
    ''' Main '''
    links = []

    try:
        if len(argv) != 3:
            raise IndexError("Not enough arguments.")

        url = argv[2].rstrip('/')
        scheme = urlparse(url).scheme

        if argv[1].upper() in ['GET', 'POST', 'OPTIONS', 'PUT', 'DELETE']:
            method = argv[1]
        else:
            raise ValueError("Method %s is not supported." % argv[1])

        r = requests.request(method, url, verify=False) # get the response
        url = r.url.rstrip('/') # if there was a redirect - acquire the URL from the response

        print('[+] Retrieving '+url)
        raw_html = r.content

        soup = BeautifulSoup(raw_html, "html.parser")

        # neat print headers
        print("[+] Received the response HTTP %d" % r.status_code)
        for header in r.headers:
            print(header + ':', r.headers[header])

        # gathering a list of hrefs and srcs
        for script in soup.find_all('script'):
            if not script == None:
                links.append(script.get('src'))
        for anchor in soup.find_all('a'):
            if not anchor == None:
                links.append(anchor.get('href'))
        for link in soup.find_all('link'):
            if not link == None:
                links.append(link.get('href'))
        for form in soup.find_all('form'):
            if not form == None:
                links.append(form.get('action'))

        links = list(filter(None, links)) # removing None src's

        links = [link for link in links if
                 not link[0] == '#' and
                 not urlparse(link).scheme in ['mailto', 'skype'] and
                 not link == '/'] # removing bookmarks, emails and skype and '/'

        links = [scheme + '://' + link[2:] if link[0:2] == '//' else
                 link if urlparse(link).scheme in ['http', 'https'] else
                 url + link if link[0] == '/' else
                 url+'/'+link for link in links] # appending URL value

    except Exception as e:
        print("[-] Something went wrong: %s" % e)

    except KeyboardInterrupt as e:
        print("[x] Exiting by user command")

    print("\n[+] %d Links extracted:" % len(links))
    for link in links:
        print(link)
    return links

if __name__ == "__main__":
    print('[*] Starting the main module')
    main(sys.argv)
\$\endgroup\$
4
\$\begingroup\$

Code style issues (from top to bottom)

  • imports - you should properly group imports as per PEP8 - in the first group you would have your system-level imports, then third-party libraries and then your local imports - all groups sorted alphabetically with a blank line between the groups:

    import sys
    from urllib.parse import urlparse
    
    from bs4 import BeautifulSoup
    import requests
    
  • docstrings should be enclosed into triple quotes, start with a capital letter and end with a dot. "Main" is though not a descriptive and useful documentation string - see if you can add details about what the function does

  • variable naming - use descriptive variable names - response instead of r
  • watch for spaces around the operators
  • when comparing to None, you should be using is and not ==, e.g. if anchor is not None - you though actually don't need these comparisons - see below
  • avoid using bare exception clauses

Other improvements

  • you should switch to argparse module instead of manually parsing and validating sys.argv
  • you don't actually need checks against None and filtering the links list - you can enforce the presence of attributes this way:

    script_elements = [element['src'] for element in soup.select('script[src]')]
    anchor_elements = [element['href'] for element in soup.select('a[href]')]
    link_elements = [element['href'] for element in soup.select('link[href]')]
    form_elements = [element['action'] for element in soup.select('form[action]')]
    
    links = script_elements + anchor_elements + link_elements + form_elements
    
  • if performance is critical, see if you can use lxml instead of html.parser (requires lxml module to be installed):

    soup = BeautifulSoup(response.content, "lxml")
    
  • see if you can join the current url with the relative link with urljoin() - might simplify the post-processing for the extracted links
\$\endgroup\$
3
  • \$\begingroup\$ much appreciated. This makes me learn the proper ways of doing things :) \$\endgroup\$ – Tom Levy Aug 11 '17 at 13:50
  • 3
    \$\begingroup\$ The requests documentation uses r as a variable name for requests.get. Is this a valid reason for using single-letter variables or would you still recommend request? \$\endgroup\$ – Daniel Aug 11 '17 at 13:52
  • 3
    \$\begingroup\$ @Coal_ good catch! It is probably not a big deal for this particular code snippet, but overall I would make a habit of making descriptive variable names unless you are doing a quick prototype. Thanks. \$\endgroup\$ – alecxe Aug 11 '17 at 14:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.