The argument that Pythagoras’s Theorem applies to any similar shapes and not just to squares appears in Euclid’s Elements, Book VI, Proposition 31. Euclid lived from about 325-265 BCE, and according to Proclus (a later Greek mathematician, 411-485 CE), this extension of Pythagoras’s Theorem is due to Euclid himself.2
For some reason, though, Euclid did not seem to be entirely happy with this proof, perhaps because of its reliance on Eudoxus’s theory of proportions, including such things as length and area scale factors. (Eudoxus lived from about 408-355 BCE.) Euclid develops the theory of proportions in Books V and VI of the Elements, but wanted to use Pythagoras’s Theorem earlier. He therefore provided an entirely different proof (which Proclus claims is due to Euclid himself3), which appears in Book I, proposition 47. You can find this proof as Approach 4 in Proving Pythagoras.