import requests
from rich.table import Table
from bs4 import BeautifulSoup
import sys
from console import console
from rich.panel import Panel
from rich.layout import Layout


layout = Layout()
layout.split_column(
    Layout(name="target"),
    Layout(name="headers"),
    Layout(name="cookies"),
    Layout(name="scripts")
)

target = sys.argv[1]

wanted_headers = [
    "Content-Security-Policy",
    "X-Frame-Option",
    "X-Content-Type-Options",
    "Strict-Transport-Security",
    "Referrer-Policy"
]

s = requests.Session()

with console.status("Making request..."):
    base_session = s.get(target)

soup = BeautifulSoup(base_session.text, "html.parser")

redirect_string = ""
if len(base_session.history) > 0:
    for history_item in base_session.history:
        redirect_string += f" -> {history_item.url} ({history_item.status_code})"
redirect_string += f" -> {base_session.url} ({base_session.status_code})"

layout["target"].split(
    Panel(target, title="Target"),
    Panel(redirect_string, title="Request Chain")
)

base_headers = base_session.headers

header_string = ""
for wanted_header in wanted_headers:
    if wanted_header not in base_headers.keys():
        header_string += f"[red]- {wanted_header}\n"
    else:
        header_string+= f"[green]- {wanted_header}\n"

layout['headers'].split(
    Panel(header_string, title="Headers")
)

if len(base_session.cookies) > 0:
    cookie_table = Table()
    cookie_table.add_column("Name", justify="right")
    cookie_table.add_column("Secure")
    cookie_table.add_column("HttpOnly")
    cookie_table.add_column("SameSite")

    for cookie in base_session.cookies:
        is_secure_string = "[green]✓" if cookie.secure else "[red]X"
        is_http_only_string = "[green]✓" if cookie.has_nonstandard_attr('HttpOnly') else "[red]X"
        is_samesite_string = f"[green]✓ ({cookie.get_nonstandard_attr('SameSite')})" if cookie.has_nonstandard_attr('SameSite') else "[red]X"

        cookie_table.add_row(cookie.name, is_secure_string, is_http_only_string, is_samesite_string)

    cookie_panel = Panel(cookie_table, title="Cookies")
else:
    cookie_panel = Panel("No Cookies", title="Cookies")

layout['cookies'].split(
    cookie_panel
)

scripts_string = ""
script_tags = soup.find_all('script')
for script in script_tags:
    if script.has_attr('src'):
        if str(script['src']).startswith("http"):
            if not script.has_attr("integrity"):
                scripts_string += f"[red]{script}\n"

layout['scripts'].split(
    Panel(scripts_string, title="Scripts")
)


console.print(layout)