Repository: ValvePython/steam

Issue:
SteamID: implement CSGO friend code
https://github.com/xPaw/SteamID.php/commit/fe934ce802f8aadc6fee1c3b47cd1a82f2328d75

Comments:

Patch:
diff --git a/steam/core/crypto.py b/steam/core/crypto.py
index 6fb2867..583bbde 100644
--- a/steam/core/crypto.py
+++ b/steam/core/crypto.py
@@ -6,7 +6,7 @@ from os import urandom as random_bytes
 from struct import pack
 from base64 import b64decode
 
-from Cryptodome.Hash import SHA1, HMAC
+from Cryptodome.Hash import MD5, SHA1, HMAC
 from Cryptodome.PublicKey.RSA import import_key as rsa_import_key, construct as rsa_construct
 from Cryptodome.Cipher import PKCS1_OAEP, PKCS1_v1_5
 from Cryptodome.Cipher import AES as AES
@@ -96,6 +96,9 @@ def hmac_sha1(secret, data):
 def sha1_hash(data):
     return SHA1.new(data).digest()
 
+def md5_hash(data):
+    return MD5.new(data).digest()
+
 def rsa_publickey(mod, exp):
     return rsa_construct((mod, exp))
 
diff --git a/steam/steamid.py b/steam/steamid.py
index 183c767..319d877 100644
--- a/steam/steamid.py
+++ b/steam/steamid.py
@@ -1,9 +1,11 @@
+import struct
 import json
 import sys
 import re
 import requests
 from steam.enums.base import SteamIntEnum
 from steam.enums import EType, EUniverse, EInstanceFlag
+from steam.core.crypto import md5_hash
 from steam.utils.web import make_requests_session
 
 if sys.version_info < (3,):
@@ -35,6 +37,7 @@ _icode_custom    = "bcdfghjkmnpqrtvw"
 _icode_all_valid = _icode_hex + _icode_custom
 _icode_map       = dict(zip(_icode_hex,    _icode_custom))
 _icode_map_inv   = dict(zip(_icode_custom, _icode_hex   ))
+_csgofrcode_chars = 'ABCDEFGHJKLMNPQRSTUVWXYZ23456789'
 
 
 class SteamID(intBase):
@@ -83,6 +86,14 @@ class SteamID(intBase):
         """
         return int(self) & 0xFFffFFff
 
+    @property
+    def account_id(self):
+        """
+        :return: account id
+        :rtype: :class:`int`
+        """
+        return int(self) & 0xFFffFFff
+
     @property
     def instance(self):
         """
@@ -197,6 +208,40 @@ class SteamID(intBase):
 
             return invite_code
 
+    @property
+    def as_csgo_friend_code(self):
+        """
+        :return: CS:GO Friend code (e.g. ``AEBJA-ABDC``)
+        :rtype: :class:`str`
+        """
+        if self.type != EType.Individual or not self.is_valid():
+            return
+
+        h = b'CSGO' + struct.pack('>L', self.account_id)
+        h, = struct.unpack('<L', md5_hash(h[::-1])[:4])
+        steamid = self.as_64
+        result = 0
+
+        for i in range(8):
+            id_nib = (steamid >> (i * 4)) & 0xF
+            hash_nib = (h >> i) & 0x1
+            a = (result << 4) | id_nib
+
+            result = ((result >> 28) << 32) | a
+            result = ((result >> 31) << 32) | ((a << 1) | hash_nib)
+
+        result, = struct.unpack('<Q', struct.pack('>Q', result))
+        code = ''
+
+        for i in range(13):
+            if i in (4, 9):
+                code += '-'
+
+            code += _csgofrcode_chars[result & 31]
+            result = result >> 5
+
+        return code[5:]
+
     @property
     def invite_url(self):
         """
@@ -413,7 +458,7 @@ def steam3_to_tuple(value):
 
     return (steam32, etype, universe, instance)
 
-def invite_code_to_tuple(code, universe=EUniverse.Public):
+def from_invite_code(code, universe=EUniverse.Public):
     """
     Invites urls can be generated at https://steamcommunity.com/my/friends/add
 
@@ -441,7 +486,45 @@ def invite_code_to_tuple(code, universe=EUniverse.Public):
     accountid = int(re.sub("["+_icode_custom+"]", repl_mapper, code), 16)
 
     if 0 < accountid < 2**32:
-        return (accountid, EType(1), EUniverse(universe), 1)
+        return SteamID(accountid, EType.Individual, EUniverse(universe), 1)
+
+SteamID.from_invite_code = staticmethod(from_invite_code)
+
+def from_csgo_friend_code(code, universe=EUniverse.Public):
+    """
+    Takes CS:GO friend code and returns SteamID
+
+    :param code: CS:GO friend code (e.g. ``AEBJA-ABDC``)
+    :type  code: :class:`str`
+    :param universe: Steam universe (default: ``Public``)
+    :type  universe: :class:`EType`
+    :return: SteamID instance
+    :rtype: :class:`.SteamID` or :class:`None`
+    """
+    if not re.match(r'^['+_csgofrcode_chars+'\-]{10}$', code):
+        return None
+
+    code = ('AAAA-' + code).replace('-', '')
+    result = 0
+
+    for i in range(13):
+        index = _csgofrcode_chars.find(code[i])
+        if index == -1:
+            return None
+        result = result | (index << 5 * i)
+
+    result, = struct.unpack('<Q', struct.pack('>Q', result))
+    accountid = 0
+
+    for i in range(8):
+        result = result >> 1
+        id_nib = result & 0xF
+        result = result >> 4
+        accountid = (accountid << 4) | id_nib
+
+    return SteamID(accountid, EType.Individual, EUniverse(universe), 1)
+
+SteamID.from_csgo_friend_code = staticmethod(from_csgo_friend_code)
 
 def steam64_from_url(url, http_timeout=30):
     """


Test patch:
diff --git a/tests/test_steamid.py b/tests/test_steamid.py
index b8e59a3..269a2cf 100644
--- a/tests/test_steamid.py
+++ b/tests/test_steamid.py
@@ -289,11 +289,22 @@ class SteamID_properties(unittest.TestCase):
                          )
 
     def test_as_invite_code(self):
-        self.assertEqual(SteamID(0     , EType.Individual, EUniverse.Public, instance=1).as_invite_code, None)
-        self.assertEqual(SteamID(123456, EType.Individual, EUniverse.Public, instance=1).as_invite_code, 'cv-dgb')
-        self.assertEqual(SteamID(123456, EType.Individual, EUniverse.Beta  , instance=1).as_invite_code, 'cv-dgb')
-        self.assertEqual(SteamID(123456, EType.Invalid   , EUniverse.Public, instance=1).as_invite_code, None)
-        self.assertEqual(SteamID(123456, EType.Clan      , EUniverse.Public, instance=1).as_invite_code, None)
+        self.assertEqual(SteamID(0         , EType.Individual, EUniverse.Public, instance=1).as_invite_code, None)
+        self.assertEqual(SteamID(1         , EType.Invalid   , EUniverse.Public, instance=1).as_invite_code, None)
+        self.assertEqual(SteamID(1         , EType.Clan      , EUniverse.Public, instance=1).as_invite_code, None)
+        self.assertEqual(SteamID(1         , EType.Individual, EUniverse.Beta  , instance=1).as_invite_code, 'c')
+        self.assertEqual(SteamID(1         , EType.Individual, EUniverse.Public, instance=1).as_invite_code, 'c')
+        self.assertEqual(SteamID(123456    , EType.Individual, EUniverse.Public, instance=1).as_invite_code, 'cv-dgb')
+        self.assertEqual(SteamID(4294967295, EType.Individual, EUniverse.Public, instance=1).as_invite_code, 'wwww-wwww')
+
+    def test_as_csgo_friend_code(self):
+        self.assertEqual(SteamID(0         , EType.Individual, EUniverse.Public, instance=1).as_csgo_friend_code, None)
+        self.assertEqual(SteamID(1         , EType.Invalid   , EUniverse.Public, instance=1).as_csgo_friend_code, None)
+        self.assertEqual(SteamID(1         , EType.Clan      , EUniverse.Public, instance=1).as_csgo_friend_code, None)
+        self.assertEqual(SteamID(1         , EType.Individual, EUniverse.Beta  , instance=1).as_csgo_friend_code, 'AJJJS-ABAA')
+        self.assertEqual(SteamID(1         , EType.Individual, EUniverse.Public, instance=1).as_csgo_friend_code, 'AJJJS-ABAA')
+        self.assertEqual(SteamID(123456    , EType.Individual, EUniverse.Public, instance=1).as_csgo_friend_code, 'ABNBT-GBDC')
+        self.assertEqual(SteamID(4294967295, EType.Individual, EUniverse.Public, instance=1).as_csgo_friend_code, 'S9ZZR-999P')
 
     def test_as_invite_url(self):
         self.assertEqual(SteamID(0     , EType.Individual, EUniverse.Public, instance=1).invite_url, None)
@@ -435,17 +446,37 @@ class steamid_functions(unittest.TestCase):
                          (1234, EType.Chat, EUniverse.Public, EInstanceFlag.Clan)
                          )
 
-    def test_arg_invite_code(self):
-        self.assertIsNone(steamid.invite_code_to_tuple('invalid_format'))
-        self.assertIsNone(steamid.invite_code_to_tuple('https://steamcommunity.com/p/cv-dgb'))
-        self.assertIsNone(steamid.invite_code_to_tuple('b'))
-        self.assertIsNone(steamid.invite_code_to_tuple('aaaaaaaaaaaaaaaaaaaaaaaaa'))
-
-        self.assertEqual(steamid.invite_code_to_tuple('cv-dgb'),
-                         (123456, EType.Individual, EUniverse.Public, 1))
-        self.assertEqual(steamid.invite_code_to_tuple('http://s.team/p/cv-dgb'),
-                         (123456, EType.Individual, EUniverse.Public, 1))
-        self.assertEqual(steamid.invite_code_to_tuple('https://s.team/p/cv-dgb'),
-                         (123456, EType.Individual, EUniverse.Public, 1))
-        self.assertEqual(steamid.invite_code_to_tuple('https://s.team/p/cv-dgb/ABCDE12354'),
-                         (123456, EType.Individual, EUniverse.Public, 1))
+    def test_from_invite_code(self):
+        self.assertIsNone(steamid.from_invite_code('invalid_format'))
+        self.assertIsNone(steamid.from_invite_code('https://steamcommunity.com/p/cv-dgb'))
+        self.assertIsNone(steamid.from_invite_code('b'))
+        self.assertIsNone(steamid.from_invite_code('aaaaaaaaaaaaaaaaaaaaaaaaa'))
+
+        self.assertEqual(steamid.from_invite_code('c', EUniverse.Beta),
+                         SteamID(1, EType.Individual, EUniverse.Beta, 1))
+        self.assertEqual(steamid.from_invite_code('c'),
+                         SteamID(1, EType.Individual, EUniverse.Public, 1))
+        self.assertEqual(steamid.from_invite_code('http://s.team/p/c', EUniverse.Beta),
+                         SteamID(1, EType.Individual, EUniverse.Beta, 1))
+        self.assertEqual(steamid.from_invite_code('http://s.team/p/c'),
+                         SteamID(1, EType.Individual, EUniverse.Public, 1))
+        self.assertEqual(steamid.from_invite_code('https://s.team/p/cv-dgb'),
+                         SteamID(123456, EType.Individual, EUniverse.Public, 1))
+        self.assertEqual(steamid.from_invite_code('https://s.team/p/cv-dgb/ABCDE12354'),
+                         SteamID(123456, EType.Individual, EUniverse.Public, 1))
+        self.assertEqual(steamid.from_invite_code('http://s.team/p/wwww-wwww'),
+                         SteamID(4294967295, EType.Individual, EUniverse.Public, 1))
+
+    def test_from_csgo_friend_code(self):
+        self.assertIsNone(steamid.from_csgo_friend_code(''))
+        self.assertIsNone(steamid.from_csgo_friend_code('aaaaaaaaaaaaaaaaaaaaaaaaaaaa'))
+        self.assertIsNone(steamid.from_csgo_friend_code('11111-1111'))
+
+        self.assertEqual(steamid.from_csgo_friend_code('AJJJS-ABAA', EUniverse.Beta),
+                         SteamID(1, EType.Individual, EUniverse.Beta, instance=1))
+        self.assertEqual(steamid.from_csgo_friend_code('AJJJS-ABAA'),
+                         SteamID(1, EType.Individual, EUniverse.Public, instance=1))
+        self.assertEqual(steamid.from_csgo_friend_code('ABNBT-GBDC'),
+                         SteamID(123456, EType.Individual, EUniverse.Public, instance=1))
+        self.assertEqual(steamid.from_csgo_friend_code('S9ZZR-999P'),
+                         SteamID(4294967295, EType.Individual, EUniverse.Public, instance=1))