import hashlib
import tempfile
import unittest
from glob import glob
from pathlib import Path

from collectionparser.crack import crack
from collectionparser.parser import Processor


class CrackTest(unittest.TestCase):
    def test_hash_saved_by_processor(self):
        with tempfile.TemporaryDirectory() as tmp:
            src = Path(tmp) / "in.txt"
            src.write_text(
                "a@x.com:5f4dcc3b5aa765d61d8327deb882cf99\nb@y.com:plainpw\n",
                encoding="utf8",
            )
            processor = Processor(
                str(Path(tmp) / "s.txt"),
                str(Path(tmp) / "f.txt"),
                str(Path(tmp) / "h.txt"),
            )
            processor.process(src)
            hashes = "".join(
                Path(p).read_text(encoding="utf8")
                for p in glob(str(Path(tmp) / "h.txt.part.*"))
            )
            success = "".join(
                Path(p).read_text(encoding="utf8")
                for p in glob(str(Path(tmp) / "s.txt.part.*"))
            )
            self.assertEqual("a@x.com:5f4dcc3b5aa765d61d8327deb882cf99\n", hashes)
            self.assertEqual("b@y.com:plainpw\n", success)

    def test_crack_md5_and_sha1(self):
        with tempfile.TemporaryDirectory() as tmp:
            hashes = Path(tmp) / "hashes.txt"
            hashes.write_bytes(
                b"a@x.com:" + hashlib.md5(b"hunter2").hexdigest().encode() + b"\n"
                b"b@y.com:" + hashlib.sha1(b"secret").hexdigest().encode() + b"\n"
                b"c@z.com:not-a-hash\n"
            )
            wordlist = Path(tmp) / "words.txt"
            wordlist.write_bytes(b"foo\nhunter2\nbar\nsecret\n")
            out = Path(tmp) / "cracked.txt"
            total = crack(str(hashes), str(out), wordlists=[str(wordlist)], workers=2)
            self.assertEqual(2, total)
            self.assertEqual(
                {b"a@x.com:hunter2", b"b@y.com:secret"},
                set(out.read_bytes().splitlines()),
            )

    def test_crack_from_plaintext_passwords(self):
        with tempfile.TemporaryDirectory() as tmp:
            hashes = Path(tmp) / "hashes.txt"
            hashes.write_bytes(
                b"a@x.com:" + hashlib.md5(b"reused").hexdigest().encode() + b"\n"
            )
            empty = Path(tmp) / "empty.txt"
            empty.write_bytes(b"")
            plaintext = Path(tmp) / "plain.txt"
            plaintext.write_bytes(b"other@x.com:reused\n")
            out = Path(tmp) / "cracked.txt"
            total = crack(
                str(hashes),
                str(out),
                wordlists=[str(empty)],
                passwords_from=[str(plaintext)],
                workers=1,
            )
            self.assertEqual(1, total)
            self.assertEqual([b"a@x.com:reused"], out.read_bytes().splitlines())

    def test_crack_chunks_hash_file(self):
        with tempfile.TemporaryDirectory() as tmp:
            hashes = Path(tmp) / "hashes.txt"
            hashes.write_bytes(
                b"a@x.com:" + hashlib.md5(b"hunter2").hexdigest().encode() + b"\n"
                b"b@y.com:" + hashlib.sha1(b"secret").hexdigest().encode() + b"\n"
            )
            wordlist = Path(tmp) / "words.txt"
            wordlist.write_bytes(b"hunter2\nsecret\n")
            out = Path(tmp) / "cracked.txt"
            total = crack(
                str(hashes),
                str(out),
                wordlists=[str(wordlist)],
                workers=1,
                max_hashes=1,
            )
            self.assertEqual(2, total)
            self.assertEqual(
                {b"a@x.com:hunter2", b"b@y.com:secret"},
                set(out.read_bytes().splitlines()),
            )


if __name__ == "__main__":
    unittest.main()