#!/usr/bin/env python3
"""
Тесты Rusgram: SLA (легальный флоу чекера) + регрессия по всем найденным дырам.

    python3 test_rusgram.py [host]

host по умолчанию — LAN-IP этой машины (важно: проверка закрытия внутренних
портов имеет смысл только с внешнего адреса; с 127.0.0.1 они доступны намеренно).

Exit code 0 — всё зелёное. 1 — есть падения.
Против services_vulnerable/rusgram секция SECURITY обязана падать,
против services/rusgram — обязана проходить.
"""
import sys
import socket
import string
import random
import requests

TIMEOUT = 6
PROXY, STATIC, DB, LOGIC, CONTENT = 38000, 38001, 38002, 38003, 38004

_results = []


def rnd(n=10):
    return "".join(random.choices(string.ascii_lowercase, k=n))


def check(section, name, passed, detail=""):
    _results.append((section, name, passed))
    mark = "\033[32mPASS\033[0m" if passed else "\033[31mFAIL\033[0m"
    print(f"  [{mark}] {name}" + (f"  ({detail})" if detail else ""), flush=True)


def lan_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        return s.getsockname()[0]
    except OSError:
        return "127.0.0.1"
    finally:
        s.close()


def port_open(host, port):
    s = socket.socket()
    s.settimeout(3)
    try:
        s.connect((host, port))
        return True
    except OSError:
        return False
    finally:
        s.close()


# ---------------------------------------------------------------- SLA
def test_sla(host):
    print("\n=== SLA: легальный флоу ===", flush=True)
    B = f"http://{host}:{PROXY}"
    s = requests.Session()
    u, p = "t_" + rnd(), rnd()
    flag = "".join(random.choices(string.ascii_uppercase + string.digits, k=31)) + "="

    try:
        r = s.post(f"{B}/api/register", json={"login": u, "password": p}, timeout=TIMEOUT)
        check("SLA", "register нового юзера = 200", r.status_code == 200, r.status_code)
        check("SLA", "выдана кука session", s.cookies.get("session") == f"{u}||{p}")

        r = s.post(f"{B}/api/send_msg", json={"msg": flag}, timeout=TIMEOUT)
        check("SLA", "send_msg = 200", r.status_code == 200, r.status_code)

        r = s.post(f"{B}/api/update_personal",
                   json={"first_name": "Ivan", "second_name": "Ivanov", "email": "i@x.ru"},
                   timeout=TIMEOUT)
        check("SLA", "update_personal = 200", r.status_code == 200, r.status_code)

        r = s.post(f"{B}/api/post_comment", json={"img_id": 3, "msg": "nice"}, timeout=TIMEOUT)
        check("SLA", "post_comment = 200", r.status_code == 200, r.status_code)

        r = s.get(f"{B}/", timeout=TIMEOUT)
        check("SLA", "GET / отдаёт владельцу его флаг", r.status_code == 200 and flag in r.text)
        check("SLA", "GET / рендерит профиль", "Ivan" in r.text)

        r = s.get(f"{B}/specific.html?img=3", timeout=TIMEOUT)
        check("SLA", "specific.html?img=3 показывает комментарий",
              r.status_code == 200 and "nice" in r.text, r.status_code)

        s2 = requests.Session()
        r = s2.post(f"{B}/api/login", json={"login": u, "password": p}, timeout=TIMEOUT)
        check("SLA", "login с верными кредами = 200", r.status_code == 200, r.status_code)
        check("SLA", "после login виден свой флаг", flag in s2.get(f"{B}/", timeout=TIMEOUT).text)

        r = requests.post(f"{B}/api/login", json={"login": u, "password": "nope"}, timeout=TIMEOUT)
        check("SLA", "login с неверными = 401", r.status_code == 401, r.status_code)

        for page in ("login.html", "register.html"):
            r = requests.get(f"{B}/{page}", timeout=TIMEOUT)
            check("SLA", f"{page} = 200", r.status_code == 200, r.status_code)
        for asset in ("css/login.css", "img/favicon.ico", "js/bootstrap.min.js"):
            r = requests.get(f"{B}/{asset}", timeout=TIMEOUT)
            check("SLA", f"статика {asset} = 200", r.status_code == 200, r.status_code)
    except requests.RequestException as e:
        check("SLA", "сервис доступен", False, str(e)[:70])
        return None, None, None
    return u, p, flag


# ---------------------------------------------------------------- SECURITY
def test_security(host, victim, vpass, vflag):
    print("\n=== SECURITY: регрессия по уязвимостям ===", flush=True)
    B = f"http://{host}:{PROXY}"

    # V1 — SSTI/RCE: полная цепочка, а не только отказ в регистрации.
    # Имя юзера с .html -> его файл рендерится content'ом как Jinja-шаблон.
    au, ap = f"{rnd()}.html", "p"
    r = requests.post(f"{B}/api/register", json={"login": au, "password": ap}, timeout=TIMEOUT)
    check("SEC", "V1: регистрация имени *.html отклонена", r.status_code != 200, r.status_code)
    canary = "SSTI" + rnd(6).upper()
    payload = "{{ lipsum.__globals__.os.popen('echo " + canary + "').read() }}"
    requests.post(f"{B}/api/update_personal",
                  json={"first_name": payload, "second_name": "x", "email": "x"},
                  cookies={"session": f"{au}||{ap}"}, timeout=TIMEOUT)
    r = requests.get(f"{B}/db/users/{au}", timeout=TIMEOUT)
    check("SEC", "V1: SSTI не выполняется (нет RCE)", canary not in r.text, r.status_code)

    # V5 — чтение файла жертвы через render_template
    r = requests.get(f"{B}/db/users/{victim}", timeout=TIMEOUT)
    check("SEC", "V1/V5: render_template произвольного пути не отдаёт файл юзера",
          vflag not in r.text, r.status_code)

    # V2 — перерегистрация чужого аккаунта
    r = requests.post(f"{B}/api/register", json={"login": victim, "password": "hacked"},
                      timeout=TIMEOUT)
    check("SEC", "V2: перерегистрация существующего юзера отклонена",
          r.status_code != 200, r.status_code)
    r = requests.get(f"{B}/", cookies={"session": f"{victim}||hacked"}, timeout=TIMEOUT)
    check("SEC", "V2: перезаписанный пароль не даёт доступа к флагу", vflag not in r.text)
    r = requests.post(f"{B}/api/login", json={"login": victim, "password": vpass}, timeout=TIMEOUT)
    check("SEC", "V2: пароль жертвы НЕ сломан (нет саботажа)", r.status_code == 200, r.status_code)

    # V3 — SSRF через Origin
    r = requests.post(f"{B}/api/send_msg", json={"msg": "x"},
                      cookies={"session": f"{victim}||WRONGPASS"},
                      headers={"Origin": f"http://database:5002/users/{victim}/msg?x="},
                      timeout=TIMEOUT)
    check("SEC", "V3: SSRF через Origin не читает внутреннюю сеть", vflag not in r.text,
          r.status_code)

    # V6 — открытый редирект
    try:
        r = requests.post(f"{B}/api/send_msg", json={"msg": "x"},
                          cookies={"session": f"{victim}||WRONGPASS"},
                          headers={"Origin": "https://evil.example"},
                          allow_redirects=False, timeout=TIMEOUT)
        loc = r.headers.get("Location", "")
        check("SEC", "V6: Location не содержит внешний домен", "evil.example" not in loc,
              loc or "нет Location")
    except requests.RequestException as e:
        check("SEC", "V6: открытый редирект", False, str(e)[:50])

    # V4/V5 — публикация внутренних портов наружу
    for port, name in ((STATIC, "static"), (DB, "database"), (LOGIC, "logic"), (CONTENT, "content")):
        closed = not port_open(host, port)
        check("SEC", f"V4/V5: порт {port} ({name}) закрыт снаружи", closed,
              "закрыт" if closed else "ОТКРЫТ")
    check("SEC", f"порт {PROXY} (proxy) открыт — сервис доступен", port_open(host, PROXY))

    # V8 — устойчивость
    cases = [
        ("GET / с кукой a||b||c", lambda: requests.get(f"{B}/", cookies={"session": "a||b||c"},
                                                       timeout=TIMEOUT)),
        ("GET / без куки", lambda: requests.get(f"{B}/", timeout=TIMEOUT)),
        ("send_msg без куки", lambda: requests.post(f"{B}/api/send_msg", json={"msg": "x"},
                                                    timeout=TIMEOUT)),
        ("send_msg с несуществующим юзером",
         lambda: requests.post(f"{B}/api/send_msg", json={"msg": "x"},
                               cookies={"session": "nosuchuser||p"}, timeout=TIMEOUT)),
        ("specific.html?img=abc",
         lambda: requests.get(f"{B}/specific.html?img=abc",
                              cookies={"session": f"{victim}||{vpass}"}, timeout=TIMEOUT)),
        ("specific.html?img=999",
         lambda: requests.get(f"{B}/specific.html?img=999",
                              cookies={"session": f"{victim}||{vpass}"}, timeout=TIMEOUT)),
        ("post_comment с img_id='../x'",
         lambda: requests.post(f"{B}/api/post_comment", json={"img_id": "../x", "msg": "y"},
                               cookies={"session": f"{victim}||{vpass}"}, timeout=TIMEOUT)),
        ("register без полей", lambda: requests.post(f"{B}/api/register", json={},
                                                     timeout=TIMEOUT)),
    ]
    for name, fn in cases:
        try:
            code = fn().status_code
            check("SEC", f"V8: {name} не даёт 500", code != 500, code)
        except requests.RequestException as e:
            check("SEC", f"V8: {name}", False, str(e)[:50])


def main():
    host = sys.argv[1] if len(sys.argv) > 1 else lan_ip()
    print(f"Цель: {host}", flush=True)
    if host.startswith("127."):
        print("  ! внутренние порты намеренно слушают 127.0.0.1 — "
              "проверка их закрытия достоверна только с внешнего IP", flush=True)

    victim, vpass, vflag = test_sla(host)
    if victim:
        test_security(host, victim, vpass, vflag)

    failed = [r for r in _results if not r[2]]
    print(f"\n{'=' * 52}")
    for sect in ("SLA", "SEC"):
        tot = [r for r in _results if r[0] == sect]
        bad = [r for r in tot if not r[2]]
        print(f"  {sect}: {len(tot) - len(bad)}/{len(tot)} passed")
    if failed:
        print(f"\n\033[31mFAILED ({len(failed)}):\033[0m")
        for _, name, _ in failed:
            print(f"    - {name}")
        return 1
    print("\n\033[32mВсё зелёное.\033[0m")
    return 0


if __name__ == "__main__":
    sys.exit(main())