SteamID: add CSGO friend code

fix #303

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))
 

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):
         """
@@ -443,6 +488,42 @@ def invite_code_to_tuple(code, universe=EUniverse.Public):
     if 0 < accountid < 2**32:
         return (accountid, EType(1), EUniverse(universe), 1)
 
+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))
+
+SteamID.from_csgo_friend_code = staticmethod(from_csgo_friend_code)
+
 def steam64_from_url(url, http_timeout=30):
     """
     Takes a Steam Community url and returns steam64 or None

tests/test_steamid.py

@@ -295,6 +295,15 @@ class SteamID_properties(unittest.TestCase):
         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)
 
+    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)
         self.assertEqual(SteamID(123456, EType.Individual, EUniverse.Public, instance=1).invite_url, 'https://s.team/p/cv-dgb')
@@ -449,3 +458,17 @@ class steamid_functions(unittest.TestCase):
                          (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_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))